zhukov000

HeatProblem

Dec 5th, 2021
617
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import numpy as np
  2. import matplotlib.pyplot as plt
  3. from math import sin, pi
  4. import time
  5. import numba
  6.  
  7. @numba.njit
  8. def q(L, dx, x, y, z, t):
  9.   center = ( (0.2 * L, 0.5 * L, 0.5 * L), (0.5 * L, 0.2 * L, 0.5 * L) )
  10.   size = ( (4 / dx, 0.2 * L, 0.2 * L), (0.2 * L, 4 / dx, 0.2 * L) )
  11.   p = (100, 100)
  12.   tetta = (50, 77)
  13.   fi = (0, 10)
  14.  
  15.   for i in range(2):
  16.     cx, cy, cz = center[i]
  17.     lx, ly, lz = size[i]
  18.  
  19.     if (cx - lx <= x <= cx + lx) and (cy - ly <= y <= cy + ly) and (cz - lz <= z <= cz + lz):
  20.       return p[i] * (sin(2 * pi * t / tetta[i] + fi[i]) + 1)
  21.  
  22.   return 0.0
  23.  
  24. @numba.njit
  25. def D(x, y, z, L):
  26.   """
  27.    Считаем коэффициент дифузии
  28.    
  29.    Если точка лежит внутри сферы или не лежит
  30.  """
  31.   if 25 * ((x - L / 2) ** 2 + (y - L / 2) ** 2 + (z - L / 2) ** 2) <= L ** 2 :
  32.     return 10**-16 # np.longdouble(10**-16)
  33.   return 1.0 # np.longdouble(1.0)
  34.  
  35. def save1DToFile(data, fileName):
  36.   """
  37.    Сохранить в файл одномерный массив:
  38.    - в первой строке размер
  39.    - далее по одному элементу в строке
  40.  """
  41.   with open(fileName, "w") as file:
  42.     file.write(str(len(data)) + '\n')
  43.     for x in data:
  44.       file.write(str(x) + '\n')
  45.  
  46. def save2DToFile(data, fileName):
  47.   """
  48.    Сохранить в файл двумерный массив:
  49.    - в первой строке размер
  50.    - далее по одному элементу в строке
  51.  """
  52.   with open(fileName, "w") as file:
  53.     file.write("%d %d\n" % (len(data), len(data[0])))
  54.     for row in data:
  55.       for x in row:
  56.         file.write(str(x) + '\n')
  57.  
  58. def save3DToFile(data, fileName):
  59.   """
  60.    Сохранить в файл трехмерный массив:
  61.    - в первой строке размер
  62.    - далее по одному элементу в строке
  63.  """
  64.   with open(fileName, "w") as file:
  65.     file.write("%d %d %d\n" % (len(data), len(data[0]), len(data[0][0])))
  66.     for m in data:
  67.       for row in m:
  68.         for x in row:
  69.           file.write(str(x) + '\n')
  70.  
  71. @numba.njit
  72. def solve(dt, dx, Tm, L, cprobe):
  73.   """
  74.    Выполнение расчетов для заданных параметров
  75.  """
  76.   iter_time = int(Tm / dt) + 1
  77.   iter_dim = int(L / dx) + 1
  78.  
  79.   T = np.empty((2, iter_dim, iter_dim, iter_dim))
  80.   T.fill(0)
  81.  
  82.   g = dt / dx ** 2
  83.  
  84.   probe = np.empty((4, iter_time))
  85.   probe.fill(0)
  86.  
  87.   cpi = [ [ int(_ / dx) for _ in cprobe[i] ] for i in range(4)] # переводим в индексы из координат
  88.  
  89.   for t in range(1, iter_time):
  90.     i = t & 1
  91.     j = 1 - i
  92.     # для всех точек кроме граней
  93.     for ix in range(1, iter_dim-1):
  94.       for iy in range(1, iter_dim-1):
  95.         for iz in range(1, iter_dim-1):
  96.           x = ix * dx
  97.           y = iy * dx
  98.           z = iz * dx
  99.  
  100.           T[i][ix][iy][iz] = (T[j][ix][iy][iz] +
  101.                            (g * D(x, y, z, L)) * (
  102.                                T[j][ix-1][iy][iz] + T[j][ix+1][iy][iz] +
  103.                                T[j][ix][iy-1][iz] + T[j][ix][iy+1][iz] +
  104.                                T[j][ix][iy][iz-1] + T[j][ix][iy][iz+1] -
  105.                                6 * T[j][ix][iy][iz] ) +
  106.                            dt * q(L, dx, x, y, z, t))
  107.  
  108.     # на гранях
  109.     # x = L
  110.     ix = iter_dim - 1
  111.     for iy in range(1, iter_dim-1):
  112.       for iz in range(1, iter_dim-1):
  113.         T[i][ix][iy][iz] = T[i][ix-1][iy][iz]
  114.    
  115.     # z = L
  116.     iz = iter_dim - 1
  117.     for iy in range(1, iter_dim-1):
  118.       for ix in range(1, iter_dim-1):
  119.         T[i][ix][iy][iz] = T[i][ix][iy][iz-1]
  120.    
  121.     # y = 0
  122.     for ix in range(1, iter_dim-1):
  123.       for iz in range(1, iter_dim-1):
  124.         T[i][ix][0][iz] = T[i][ix][1][iz]
  125.    
  126.     # z = 0
  127.     for ix in range(1, iter_dim-1):
  128.       for iy in range(1, iter_dim-1):
  129.         T[i][ix][iy][0] = T[i][ix][iy][1]
  130.  
  131.     # пробы
  132.     for k in range(4):
  133.       ix, iy, iz = cpi[k]
  134.       probe[k][t] = T[i][ix][iy][iz]
  135.  
  136.   return probe, T[ 1 - (iter_time & 1) ]
  137.  
  138.  
  139. def printColorMap(fileName, L, dx, cprobe):
  140.   """
  141.    Цветовая карта в сечении 0.5 * L в конечный момент времени
  142.  """
  143.  
  144.   z = int(0.5 * L / dx)
  145.   # считываем данные из файла
  146.   with open(fileName) as file:
  147.     X, Y, Z = map(int, file.readline().split())
  148.     data = [ [ 0 for j in range(Y) ] for i in range(X) ]
  149.  
  150.     for ix in range(X):
  151.       for iy in range(Y):
  152.         for iz in range(Z):
  153.           v = float(file.readline())
  154.           if iz == z:
  155.             data[ix][iy] = v
  156.  
  157.   # data = data[::-1][1:]
  158.   data = list(zip(*data))[::-1]
  159.   Y -= 1
  160.   # подготавливаем карту
  161.   fig, ax = plt.subplots(figsize=(40, 8))
  162.  
  163.   plt.imshow(data, cmap=plt.get_cmap('coolwarm'), extent=[0, X, 0, Y])
  164.   ax.set_aspect(2)
  165.  
  166.   stepX = max(1, int(4 * (X - 1) / L)) * 4
  167.   stepY = max(1, int(2 * (X - 1) / L)) * 2
  168.  
  169.   ticksForX = [_ for _ in range(0, X, stepX)]
  170.   ticksLabelForX = ["{:1.2f}".format(_ * dx) for _ in range(0, X, stepX)]
  171.  
  172.   ticksForY = [_ for _ in range(0, X, stepY)]
  173.   ticksLabelForY = ["{:1.2f}".format(_ * dx) for _ in range(0, X, stepY)]
  174.  
  175.   ax.set_yticks(ticksForY)
  176.   ax.set_yticklabels(ticksLabelForY)
  177.   ax.set_xticks(ticksForX)
  178.   ax.set_xticklabels(ticksLabelForX)
  179.  
  180.   ax.set_xlabel("ось X")
  181.   ax.set_ylabel("ось Y")
  182.  
  183.   # пробы в виде точек на цветовой карте
  184.   for i in range(len(cprobe)):
  185.     x, y, z = cprobe[i]
  186.     ix, iy = int(x/dx), int(y/dx)
  187.     ax.plot([ix], [iy], 'o', color='black')
  188.     ax.annotate(f'({i})', xy=(ix, iy), xytext=(ix+1, iy))
  189.  
  190.   plt.colorbar()
  191.  
  192.  
  193. def loadProbesFromFiles(probeFile):
  194.   """
  195.    Читаем данные с температурами в пробах из файла для заданного случая case
  196.  """
  197.   probes = [ ]
  198.   for i in range(4):
  199.     with open(probeFile % (i)) as file:
  200.       n = int(file.readline())
  201.       probes.append([])
  202.       for _ in file:
  203.         probes[i].append(float(_))
  204.   return probes
  205.  
  206. def showProbes(probeFile, Tm, L, dx, dt):
  207.   """
  208.    Отобразить график изменения температуры в пробах
  209.  """
  210.   # считываем данные из файлов
  211.   probes = loadProbesFromFiles(probeFile)
  212.   n = len(probes[0])
  213.  
  214.   # рисуем графики
  215.   fig, ax = plt.subplots(figsize=(10, 10))
  216.   t = [_ for _ in range(n)]
  217.   for i in range(4):  
  218.     plt.plot(t, probes[i], label=f'проба{i}')
  219.  
  220.   ticksForX = [_ for _ in range(0, n, (n+9)//10)]
  221.   ticksLabelForX = ["{:1.2f}".format(_ * dt) for _ in range(0, n, (n+9)//10)]
  222.   ax.set_xticks(ticksForX)
  223.   ax.set_xticklabels(ticksLabelForX)
  224.   plt.xlabel("Ось времени (c)")
  225.  
  226.   plt.ylabel("Температура")
  227.  
  228.   plt.legend()
  229.   plt.show()
  230.  
  231.  
  232. def doFFT(probeFile, Tm, dt):
  233.   """
  234.    Восстановление периодов колебаний с использованием быстрого преобразования Фурье
  235.  """
  236.   from scipy.signal import find_peaks
  237.  
  238.   # считываем данные из файлов
  239.   probes = loadProbesFromFiles(probeFile)
  240.  
  241.   # fig, ax = plt.subplots(figsize=(10, 10))
  242.   for i in range(4):  
  243.     sig = np.array(probes[i], dtype=float)
  244.     sig_fft = np.fft.rfft(sig) # строим быстрое преобразование Фурье
  245.     freq = np.fft.rfftfreq(sig.size, d=dt)
  246.    
  247.     peaks, _ = find_peaks(np.abs(sig_fft))
  248.    
  249.     print(f'проба{i}, периоды:', end=' ')
  250.     for p in peaks:
  251.       print(f'{1 / freq[p] / dt}', end=' ')
  252.     print()
  253.  
  254.  
  255. def calculateAll(params, L, Tm, cprobe):
  256.   """
  257.    Посчитать все случаи заданные в `params` и сохранить данные в файлы
  258.    или загрузить данные из файлов для отображения
  259.  """
  260.   import os.path
  261.  
  262.   cp = np.empty((len(cprobe), 3))
  263.   for i in range(len(cprobe)):
  264.     for j in range(3):
  265.       cp[i][j] = cprobe[i][j]
  266.  
  267.   for i in range(len(params)):
  268.     dx, dt = params[i]
  269.  
  270.     # проверить, если нет файла расчета - то выполнить
  271.     # dirPath = 'sample_data/save_data/'
  272.     heatFieldFile = dirPath + f'heat-field-case{i+1}.txt'
  273.     probeFile = dirPath + f"probe%d-case{i+1}.txt"
  274.     timeFile = dirPath + f"time-case{i+1}.txt"
  275.  
  276.     if not os.path.isfile(timeFile):
  277.       startTime = time.time()
  278.       probe, T = solve(dt, dx, Tm, L, cp)
  279.       endTime = time.time()
  280.       # сохранение в файл времени расчета
  281.       calcTime = f"dx = {dx},  dt = {dt}, calculation time = {endTime-startTime} с."
  282.       with open(timeFile, "w") as file:
  283.         file.write(calcTime)
  284.       # сохранение в файлы температуры в пробах
  285.       for j in range(len(probe)):
  286.         save1DToFile(probe[j], probeFile % (j))
  287.       # сохранение в файл температурного поля в конечный момент времени
  288.       save3DToFile(T, heatFieldFile)
  289.     else:
  290.       # если файлы расчетов есть, то не выполнять расчет
  291.       # вывод времени расчета для предыдущего случая
  292.       with open(timeFile) as file:
  293.         calcTime = file.readline().strip()
  294.    
  295.     # вывод времени расчета
  296.     print(calcTime)      
  297.  
  298.     # построение цветовой карты
  299.     printColorMap(heatFieldFile, L, dx, cprobe)
  300.    
  301.     # построение графиков проб
  302.     showProbes(probeFile, Tm, L, dx, dt)
  303.    
  304.     # восстановление периодов
  305.     doFFT(probeFile, Tm, dt)
  306.  
  307.  
  308. params = ((4, 0.2), (2, 0.1), (1, 0.05), (0.5, 0.025))
  309. L = 100
  310. Tm = 600
  311. cprobe = ( (0.3 * L, 0.3 * L, 0.5 * L), (0.3 * L, 0.7 * L, 0.5 * L), (0.7 * L, 0.7 * L, 0.5 * L), (0.7 * L, 0.3 * L, 0.5 * L) )
  312. calculateAll(params, L, Tm, cprobe)
RAW Paste Data