Advertisement
zhukov000

HeatProblem

Dec 5th, 2021
700
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 9.98 KB | None | 0 0
  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)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement