{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "97c868dd-b880-4af1-af59-7dc5663861fa",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Numeryczne rozwiązywanie równania dyfuzji"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e11ea57-8816-446f-95ca-b516eb0a0754",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Printf\n",
    "using GLMakie\n",
    "using LinearAlgebra\n",
    "using LaTeXStrings\n",
    "\n",
    "my_theme = Theme(\n",
    "               Axis = (\n",
    "                   xlabelsize = 28, \n",
    "                   xticklabelsize = 24, \n",
    "                   xminorticks = IntervalsBetween(5), \n",
    "                   xminorticksvisible = true, \n",
    "                   xgridvisible = false,\n",
    "                   ylabelsize = 28, \n",
    "                   yticklabelsize = 24, \n",
    "                   ygridvisible = false, \n",
    "                   yminorticks = IntervalsBetween(5), \n",
    "                   yminorticksvisible = true,\n",
    "               ),\n",
    "               Legend = (\n",
    "                   labelsize = 24,\n",
    "               ),\n",
    "                   palette=(color = cgrad(:managua25, 25, rev=true),),\n",
    "                   Lines = (cycle = [:color],)\n",
    "            )\n",
    "set_theme!(my_theme)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e643d0e5-6899-410a-ac40-e6b3b9840636",
   "metadata": {},
   "source": [
    "Nasz reaktor będzie składał się z komórek, dla których musimy podać parametry materiału: stałą dyfuzji, przekrój czynny na absporpcję i na rozszczepienie.\n",
    "Wygodnie będzie zdefiniować strukturę zawierającą takie informacje. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5985eee8-aa34-485d-9c80-14b79b59cd85",
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Material\n",
    "    D::Float64\n",
    "    Σa::Float64\n",
    "    νΣf::Float64\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbfebb9d-2f62-4daf-a0db-54d576685b91",
   "metadata": {},
   "source": [
    "## Reaktor 1D"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61fc8ab9-489d-4a68-9633-a25609259d9b",
   "metadata": {},
   "source": [
    "Spróbujmy rozwiązać jednowymiarowy problem reaktora taflowego otoczonego reflektorem.\n",
    "\n",
    "Zacznijmy od zdefiniowania kilku materiałów. Będziemy potrzebować rdzenia, reflektora (np. wody) oraz próżni. Parametry tego typu materiałów pojawiały się na zajęciach i pożyczymy wartości z poprzednich zadań. Natomiast próżnia będzie miała o zerowe przekroje czynne na absorpcję i rozszczepienie oraz bardzo dużą stałą dyfuzji - jest ona odwrotnością przekroju na rozpraszanie, więc dla próżni powinna dążyć do nieskończoności."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b2c99a8-c88f-4e24-9903-c0b1c359a2c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "core = Material(9.21, 0.153, 0.157)\n",
    "moderator = Material(6.4, 0.0196, 0.0)\n",
    "void = Material(1.0e6, 0.0, 0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87973346-9c51-4a66-b245-d4756afa0582",
   "metadata": {},
   "source": [
    "Załóżmy, że grubość reaktora będzie równa $2R_c$, a grubość reflektora $R_r$, z każdej strony (jednostki będą umowne). Układ dodatkowo otoczymy $R_a$ komórkami próżni, tak, aby warunek brzegowy mógł być łagodnie spełniony."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3b9e301-3177-4126-8e36-5290d84fc9df",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ra = 2\n",
    "Rr = 5\n",
    "Rc = 25\n",
    "N = 2 * (Ra + Rc + Rr)\n",
    "\n",
    "reactor = Array{Material, 1}(undef, N)\n",
    "for i in 1:Ra\n",
    "    reactor[i] = void\n",
    "end\n",
    "for i in Ra+1:Ra+Rr\n",
    "    reactor[i] = moderator\n",
    "end\n",
    "for i in Ra+Rr+1:Ra+Rr+2*Rc\n",
    "    reactor[i] = core\n",
    "end\n",
    "for i in Ra+Rr+2*Rc+1:Ra+Rr+2*Rc+Rr\n",
    "    reactor[i] = moderator\n",
    "end\n",
    "for i in Ra+Rr+2*Rc+Rr+1:N\n",
    "    reactor[i] = void\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1230e6e5-3e40-4704-9c2a-52016aea1064",
   "metadata": {},
   "source": [
    "Sprawdźmy jak wygląda nasz reaktor rysując przekrój czynny na absorpcję."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1796b8d1-b3ef-42ee-97f6-e538a48f1f54",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = Figure(size=(800, 500))\n",
    "ax = Axis(fig[1, 1]; xlabel=L\"r\", ylabel=L\"\\Sigma_a\", limits=(0, nothing, 0, nothing))\n",
    "stairs!(ax, [reactor[i].Σa for i in 1:N], color=:tomato, linewidth=2)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c96d83bb-bd81-4695-ac6a-b1b4a3b8d1e7",
   "metadata": {},
   "source": [
    "Główna część problemu polega na zaimplementowaniu metody rozwiązywania zagadnienia danego przez macierz `reactor`. \n",
    "Będziemy musieli przejść przez kilka kroków, które pokażemy osobno, ale finalnie zostaną zebrane w jedną funkcję.\n",
    "\n",
    "Rozwiązaniem naszego problemu będzie strumień neutronów $\\Phi$. Musimy zacząć od jakiegoś rozwiązania początkowego, które powinno być w miarę rozsądne (np. wektor samych zer jest złym pomysłem). Im lepiej zgadniemy postać rozwiązania, tym szybciej algorytm zbiegnie do prawidłowego. Prosty pomysł spełniający te warunki to np. wektor zawierający stałą tam, gdzie mamy jakiś materiał rozszczepialny i zera w pozostałych miejscach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8e906ed-f019-4079-a109-ed678f9f1ec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = zeros(N)\n",
    "for i in 1:N\n",
    "    if reactor[i].νΣf > 0.0\n",
    "        F[i] = 1.0\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b171d9b-bf73-4164-b101-3ebb561fb208",
   "metadata": {},
   "source": [
    "Następny krok to stworzenie macierzy A, o rozmiarze $N\\times N$, definującej nasze zdyskretyzowane równiania dyfuzji dla każdego punktu reaktora.\n",
    "\n",
    "Implementujemy tutaj schemat sprzężeń dla problemu 1D omawiany na ćwiczeniach.\n",
    "\n",
    "Uwaga o wyrazach A[1] i A[N] - w przypadku warunku brzegowego \"próżni\", ze względu na brak sąsiadów na zewnątrz, musimy zadbać o brak wpływów z tej strony. Alternatywą jest warunek brzegowy cykliczny (stosowany np. w modelowaniu nieskończonych siatek), wtedy możemy zażądać, aby A[1] było sąsiadem A[N] i odpowiednio zmodyfikować równania."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "471a0bae-d75e-47ee-8f01-ebad3667ea80",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = zeros(N, N)\n",
    "d = 1.0\n",
    "for i in 1:N\n",
    "    A[i, i] = reactor[i].Σa + 2 * reactor[i].D / (2 * d^2)\n",
    "    if i > 1\n",
    "        A[i, i] += reactor[i-1].D / (2 * d^2)\n",
    "        A[i, i-1] = -(reactor[i].D + reactor[i-1].D) / (2 * d^2)\n",
    "    end\n",
    "    if i < N\n",
    "        A[i, i] += reactor[i+1].D / (2 * d^2)\n",
    "        A[i, i+1] = -(reactor[i].D + reactor[i+1].D) / (2 * d^2)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4fbfc9d-2a79-4d3f-aebd-c526f9d1651d",
   "metadata": {},
   "source": [
    "Sprawdźmy, jak wygląda macierz A."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c817b83-1344-4132-8ae0-f47453c2696a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = Figure(size=(600, 500))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"i\", ylabel=\"j\")\n",
    "hm = heatmap!(ax, replace(A, 0.0=>NaN), colormap=:lighttest, colorrange=(-20, 20))\n",
    "Colorbar(fig[1, 2], hm)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "213e339b-d85c-417d-9c00-995dbf66048d",
   "metadata": {},
   "source": [
    "Potrzebujemy teraz parametru zbieżności rozwiązania ($\\lambda$), oraz $\\epsilon$ względną różnicę między nową $\\lambda_{i+1}$, a poprzednią $\\lambda_i$. Ustalmy, że uznajemy rozwiązanie, za zbieżne, gdy epsilon jest mniejszy niż $10^{-6}$.\n",
    "\n",
    "Na podstawie parametrów reaktora, lambdy oraz wstępnego strumienia policzymy prawą część równania (źródło: $S$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b92fe1e0-8cbe-4fb0-b741-30e9dd8946d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "λ = 1.0\n",
    "S = 1 / λ .* [reactor[i].νΣf for i in 1:N] .* F\n",
    "ϵ = 1.0\n",
    "ϵlim = 1.0e-6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dff3a4c8-9290-42f9-9a40-b6e19bb1b734",
   "metadata": {},
   "source": [
    "Jesteśmy gotowi do stworzenia głównej pętli rozwiązania. Dopóki $\\epsilon$ - względna zmiana $\\lambda_i$ - jest większy niż zadana granica, będziemy powtarzać schemat:\n",
    "1. Znajdź nowy strumień $\\Phi = A^{-1}S$\n",
    "1. Policz nowe $\\lambda_{i+1}$ (suma po wszystkich punktach w prawej części równania z nowym strumieniem, podzielona przez poprzednią wartość)\n",
    "1. Policz nową prawą cześć równania (dla $\\lambda_{i+1}$)\n",
    "1. Policz względną zmianę lambdy\n",
    "\n",
    "W Julii do rozwiązania zagadnienia $A\\Phi = S$ użyjemy operatora dzielenia macierzowego `\\`, który dla operacji A \\ B zwraca wynik X, taki, że A * X = B, czyli to o co nam chodzi. Operator ten jest polialgorytmiczny - w zależności od charakteru macierzy sam dobiera właściwą metodę. Dla macierzy górnotrójkątnych stosuje podstawienie wsteczne, dla trójdiagonalnych rozkład LU. O tym jak rozpozna naszą macierz możemy przekonać się wołając funkcję `factorize`, która zwróci najwygodniejszą w danym przypadku postać."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeec4b37-16b2-464d-b99e-df391a32a6a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "factorize(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f96c4679-11c8-4296-8336-d5eada9b6d54",
   "metadata": {},
   "outputs": [],
   "source": [
    "Sf = [reactor[i].νΣf for i in 1:N]\n",
    "while ϵ > ϵlim\n",
    "    Fn = A \\ S\n",
    "    λₙ = sum(Sf .* Fn) / sum(Sf .* F) * λ\n",
    "    Sn = 1 / λₙ .* Sf .* Fn\n",
    "    ϵ = abs((λₙ - λ) / λ)\n",
    "    F = copy(Fn)\n",
    "    S = copy(Sn)\n",
    "    λ = λₙ\n",
    "    @printf(\"ϵ = %.3e\\n\", ϵ)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "118e64c2-9304-4db9-84b7-5c34f29e7a7d",
   "metadata": {},
   "source": [
    "Dla wygody wektor zawierający wyraz $\\nu\\Sigma_f$ policzyliśmy przed główną pętlą, aby nie powtarzać zbędnie operacji zbierania tej struktury.\n",
    "\n",
    "Narysujmy uzyskany strumień."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0efa6c10-c4a6-40f6-a61a-7bfb5c25d654",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = Figure(size=(800, 600))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"r\", ylabel=L\"\\Phi\")\n",
    "lines!(ax, F, color=:tomato, linewidth=2)\n",
    "vlines!(ax, [Ra, Ra+Rr, Ra+Rr+2*Rc, Ra+Rr+2*Rc+Rr], color=:black, linestyle=:dash)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "757c12dd-0192-4fbb-bab3-a60257be9e29",
   "metadata": {},
   "source": [
    "Możemy sprawdzić też jak wyglądają kolejne przybliżenia, rysując strumień w pętli dla kolejnych przybliżeń."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0406e9b2-d1ad-4d26-babc-b97bcb39df03",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = zeros(N)\n",
    "for i in 1:N\n",
    "    if reactor[i].νΣf > 0.0\n",
    "        F[i] = 1.0\n",
    "    end\n",
    "end\n",
    "\n",
    "λ = 1.0\n",
    "S = 1 / λ .* [reactor[i].νΣf for i in 1:N] .* F\n",
    "ϵ = 1.0\n",
    "ϵlim = 1.0e-6\n",
    "i = 1\n",
    "Sf = [reactor[i].νΣf for i in 1:N]\n",
    "\n",
    "fig = Figure(size=(800, 500))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"r\", ylabel=L\"\\Phi\")\n",
    "vlines!(ax, [Ra, Ra+Rr, Ra+Rr+2*Rc, Ra+Rr+2*Rc+Rr], color=:black, linestyle=:dash)\n",
    "\n",
    "while ϵ > ϵlim\n",
    "    lines!(ax, F)\n",
    "    Fn = A \\ S\n",
    "    λₙ = sum(Sf .* Fn) / sum(Sf .* F) * λ\n",
    "    Sn = 1 / λₙ .* Sf .* Fn\n",
    "    ϵ = abs(λₙ - λ) / λ\n",
    "    F = copy(Fn)\n",
    "    S = copy(Sn)\n",
    "    λ = λₙ\n",
    "    @printf(\"i = %d, ϵ = %.3e\\n\", i, ϵ)\n",
    "    i += 1\n",
    "\n",
    "end\n",
    "lines!(ax, F)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f2311a8-2e37-4bd5-9eb6-259eecc160a2",
   "metadata": {},
   "source": [
    "Wreszcie naszą metodę można zebrać do funkcji, która będzie uniwersalna dla problemów 1D."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99263940-b907-4470-99d8-b5b9c94a61c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "function solve1d(reactor; ϵlim=1e-6)\n",
    "    N = length(reactor)\n",
    "    F = zeros(N)\n",
    "    for i in 1:N\n",
    "        if reactor[i].νΣf >= 0.0\n",
    "            F[i] = 1.0\n",
    "        end\n",
    "    end\n",
    "\n",
    "    λ = 1.0\n",
    "    S = 1 / λ .* [reactor[i].νΣf for i in 1:N] .* F\n",
    "    ϵ = 1.0\n",
    "    Sf = [reactor[i].νΣf for i in 1:N]\n",
    "\n",
    "    while ϵ > ϵlim\n",
    "        Fn = A \\ S\n",
    "        λn = sum(Sf .* Fn) / sum(Sf .* F) * λ\n",
    "        Sn = 1 / λn .* Sf .* Fn\n",
    "        ϵ = abs(λn - λ) / λ\n",
    "        F = copy(Fn)\n",
    "        S = copy(Sn)\n",
    "        λ = λn\n",
    "    end\n",
    "    F\n",
    "end\n",
    "\n",
    "F = solve1d(reactor)\n",
    "lines(F, color=:tomato)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "538f5978-b034-49d2-8e44-74275061f4d3",
   "metadata": {},
   "source": [
    "## Reaktor 2D"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ba463e9-5632-4649-baec-05a7d26aec44",
   "metadata": {},
   "source": [
    "W przypadku reaktora dwuwymiarowego nasze komórki ustawiamy w wektorze jednowymiarowym (odpowiednio definiując indeks), \n",
    "tak aby macierz A nadal była dwuwymiarowa - z tym, że jej wymiar będzie teraz $N^2 \\times N^2$. Najwygodniej będzie stworzyć funkcję obliczającą taki indeks z danych położenia komórki $(i, j)$.\n",
    "\n",
    "Będziemy musimy zmodyfikować macierz opisującą problem - pojawią się wyrazy dające sprzężenia z sąsiadami \"z góry\" i \"z dołu\", oprócz prawego i lewego, których mieliśmy w zagadnieniu 1D.\n",
    "\n",
    "Uwaga - Julia stosuje porządek macierzy kolumnowy i nasze indeksowanie ustawimy zgodnie z tą logiką. Podobnie, iteracja po pętlach powinna najpierw używać indeksu kolumny, a potem wiersza. W językach jak C czy Python, stosowany jest porządek wierszowy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95035886-b24d-4d4c-8ce5-723c529d0afa",
   "metadata": {},
   "outputs": [],
   "source": [
    "reshape(1:9, 3, 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8551d05e-db43-44c3-a71e-07c615f85137",
   "metadata": {},
   "source": [
    "Pierwszy indeks to rząd, drugi kolumna, a zatem funkcja zwracająca indeks będzie miała postać następującą (pod spodem test naszego indeksowania)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da4da26e-a2d7-4bad-b81b-5947faaeaf22",
   "metadata": {},
   "outputs": [],
   "source": [
    "function k_index(i, j, N)\n",
    "    i + (j - 1) * N\n",
    "end\n",
    "\n",
    "for j in 1:3\n",
    "    for i in 1:3\n",
    "        @printf(\"(%d,%d)->%d  \", i, j, k_index(i, j, 3))\n",
    "    end\n",
    "    println()\n",
    "end\n",
    "println()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f0ee98b-a98e-473c-90a2-8d561ee2e163",
   "metadata": {},
   "source": [
    "Zaczniemy od stworzenia reaktora. Ponownie spróbujemy rozwiązać problem z reflektorem, tym razem jednak reaktor będzie miał przekrój kołowy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4bec0c4-93f5-490f-bdbd-fe4bf828288f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ra = 2\n",
    "Rr = 5\n",
    "Rc = 25\n",
    "N = 2 * (Ra + Rr + Rc + 1)\n",
    "\n",
    "reactor2 = Array{Material, 2}(undef, N, N)\n",
    "for j in 1:N, i in 1:N\n",
    "    r = sqrt((i - N/2)^2 + (j - N/2)^2)\n",
    "    if r > Rr + Rc\n",
    "        reactor2[i, j] = void\n",
    "    elseif Rr + Rc >= r > Rc\n",
    "        reactor2[i, j] = moderator\n",
    "    else \n",
    "        reactor2[i, j] = core\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5877b42-61ad-46de-9d61-d07139c115c4",
   "metadata": {},
   "source": [
    "Znowu sprawdzimy jak wygląda nasz reaktor patrząc na $\\Sigma_a$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6453403-0d29-44c1-996c-835cf403719c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = Figure(size=(800, 600))\n",
    "Sa = zeros(N, N)\n",
    "for j in 1:N, i in 1:N\n",
    "    Sa[i, j] = reactor2[i, j].Σa\n",
    "end\n",
    "ax = Axis(fig[1, 1]; xlabel=\"x\", ylabel=\"y\")\n",
    "hm = heatmap!(ax, Sa, colormap=:lighttest)\n",
    "Colorbar(fig[1, 2], hm)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25c4eeee-b328-4f8a-9c5d-a4bd72eee1a4",
   "metadata": {},
   "source": [
    "Rozwiązywanie znowu zaczniemy od wstępnego rozwiązania, które jak poprzednio będzie zawierać jedynki dla obszaru rdzenia i zera dla pozostałych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a36599f-6c47-4a49-af04-0b517d1d7305",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = zeros(N, N)\n",
    "for j in 1:N, i in 1:N\n",
    "    if reactor2[i, j].νΣf > 0\n",
    "        F[i, j] = 1.0\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f1245f1-3576-4500-90c0-a1810f795bb2",
   "metadata": {},
   "source": [
    "Tworzymy teraz macierz A, która będzie definiować sprzężenia między komórkami. Ponieważ komórek jest $N^2$, a macierz musi podawać współczynniki dla dowolnej pary, będzie miała rozmiar $N^2\\times N^2$. \n",
    "\n",
    "Wypełniamy ją wartościami zgodnie ze schematem dla problemów 2D. Dla każdej komórki $k$ liczymy indeks sąsiadów góra/dół/lewo/prawo - odpowiednio $k_u, k_d, k_l, k_r$. Tak jak poprzednio musimy zadbać o przypadki bez danych sąsiadów."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db10999a-91f0-4d15-a691-972c1d62c267",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = zeros(N^2, N^2)\n",
    "for j in 1:N, i in 1:N\n",
    "    k = k_index(i, j, N)\n",
    "    ku = k_index(i - 1, j, N)\n",
    "    kd = k_index(i + 1, j, N)\n",
    "    kl = k_index(i, j - 1, N)\n",
    "    kr = k_index(i, j + 1, N)\n",
    "    A[k, k] = reactor2[k].Σa + 2 * reactor2[k].D / d^2\n",
    "    if j > 1\n",
    "        A[k, kl] -= (reactor2[k].D + reactor2[kl].D) / (2 * d^2)\n",
    "        A[k, k] += reactor2[kl].D / (2 * d^2)\n",
    "    end\n",
    "    if j < N\n",
    "        A[k, kr] -= (reactor2[k].D + reactor2[kr].D) / (2 * d^2)\n",
    "        A[k, k] += reactor2[kr].D / (2 * d^2)\n",
    "    end\n",
    "    if i > 1\n",
    "        A[k, ku] -= (reactor2[k].D + reactor2[ku].D) / (2 * d^2)\n",
    "        A[k, k] += reactor2[ku].D / (2 * d^2)\n",
    "    end\n",
    "    if i < N\n",
    "        A[k, kd] -= (reactor2[k].D + reactor2[kd].D) / (2 * d^2)\n",
    "        A[k, k] += reactor2[kd].D / (2 * d^2)\n",
    "    end\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3f7cc06-af21-4995-b96d-35ee50088754",
   "metadata": {},
   "source": [
    "Sprawdzamy postać macierzy A, ponieważ jest ona dość duża i w większości pusta, po prawiej stronie\n",
    "powiększenie obszaru 1-100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b81699-174d-4894-a437-1d3276833022",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = Figure(size=(1200, 600))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"i\", ylabel=\"j\")\n",
    "hm = heatmap!(ax, replace(A, 0.0=>NaN), colormap=:lighttest, colorrange=(-10, 10))\n",
    "Colorbar(fig[1, 2], hm)\n",
    "\n",
    "ax2 = Axis(fig[1, 3]; xlabel=\"i\", ylabel=\"j\")\n",
    "heatmap!(ax2, replace(A, 0.0=>NaN)[1:100, 1:100], colormap=:lighttest, colorrange=(-10, 10))\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b39c455-e19f-479d-a191-5df74672cbb9",
   "metadata": {},
   "source": [
    "Z ciekawości możemy sprawdzić jaki algorytm będzie zastosowany do tej macierzy, która ma teraz 5 diagonalii."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574422ca-1797-4a8e-af7a-90a5b05830a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "factorize(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab8f1d64-5ebb-4fd2-8ec5-d608d4b31f81",
   "metadata": {},
   "source": [
    "Algorytm wybrał rozkład Buncha-Kaufmana $A=PLDL'P'$ dla macierzy kwadratowej i symetrycznej. W poprzedniej wersji biblioteki wybierał\n",
    "zbliżony rozkład Choleskiego ($A = LL^T$), który jest możliwy dla macierzy symetrycznej i dodatnio określonej. Jest on około dwa razy bardziej wydajny niż jeszcze ogólniejszy przypadek rozkładu LU. Metoda B-K jest wydajniejsza dla macierzy blokowych i rzadkich."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72803c05-b35c-458e-8c8a-ab0dd0ddc39f",
   "metadata": {},
   "source": [
    "Jeżeli nasz początkowy strumień zamienimy na jednowymiarowy wektor, to zasadniczo pozostała część rozwiązania numerycznego powinna działać tak samo - definicja zagadnienia kryje się w macierzy A."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8289198-b2fb-4030-b689-260d8b6297f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = reshape(F, N^2)\n",
    "\n",
    "Sf = [reactor2[i].νΣf for i in eachindex(reactor2)]\n",
    "λ = 1.0\n",
    "S = 1 / λ .* Sf .* F\n",
    "ϵ = 1.0\n",
    "\n",
    "while ϵ > ϵlim\n",
    "    Fn = A \\ S\n",
    "    λn = sum(Sf .* Fn) / sum(Sf .* F) * λ\n",
    "    Sn = 1 / λn .* Sf .* Fn\n",
    "    ϵ = abs(λn - λ) / λ\n",
    "    F = copy(Fn)\n",
    "    S = copy(Sn)\n",
    "    λ = λn\n",
    "end\n",
    "@printf(\"%.3e\\n\", ϵ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "669cc883-257d-4f2d-8179-531aee722027",
   "metadata": {},
   "source": [
    "Jeżeli chcemy narysować mapę strumienia wystarczy po rozwiązaniu wrócić do postaci $N\\times N$. Warto zwrócić uwagę, że reinterpretacja macierzy do innego kształtu zasadniczo nic nie kosztuje - elementy macierzy są i tak ułożone w pamięci w sposób liniowy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50680556-4a64-4c5b-8856-970a83997edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "F = reshape(F, N, N)\n",
    "fig = Figure(size=(800, 600))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"x\", ylabel=\"y\")\n",
    "hm = heatmap!(ax, F, colormap=:YlGn)\n",
    "Colorbar(fig[1, 2], hm)\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11ee43e7-7fee-4b6b-a4ac-fb4891e18802",
   "metadata": {},
   "source": [
    "Aby lepiej zrozumieć kształt rozwiązania, narysujemy przekroje wzdłuż wybranych osi X $(0, R_c/2, Rc)$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cad96a60-b828-4116-9b7f-6fdd2ced8391",
   "metadata": {},
   "outputs": [],
   "source": [
    "x0 = round(Int, N/2)\n",
    "fig = Figure(size=(800, 600))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"y\", ylabel=L\"\\Phi\")\n",
    "lines!(ax, F[x0, :], label=L\"0\")\n",
    "lines!(ax, F[x0-round(Int, Rc/2), :], label=L\"R_c / 2\")\n",
    "lines!(ax, F[x0-Rc, :], label=L\"R_c\")\n",
    "axislegend(ax)\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b812db9b-b5d8-4939-812f-00ae416baf93",
   "metadata": {},
   "source": [
    "Jak poprzednio naszą metodę możemy zebrać do jednej funkcji."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60d631b3-366d-4343-b169-d830b801db0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "function solve2d(reactor2; ϵlim=1.0e-6)\n",
    "    N = size(reactor2)[1]\n",
    "    F = zeros(N, N)\n",
    "    for j in 1:N, i in 1:N\n",
    "        if reactor2[i, j].νΣf > 0\n",
    "            F[i, j] = 1.0\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    A = zeros(N^2, N^2)\n",
    "    for j in 1:N, i in 1:N\n",
    "        k = k_index(i, j, N)\n",
    "        ku = k_index(i - 1, j, N)\n",
    "        kd = k_index(i + 1, j, N)\n",
    "        kl = k_index(i, j - 1, N)\n",
    "        kr = k_index(i, j + 1, N)\n",
    "        A[k, k] = reactor2[k].Σa + 2 * reactor2[k].D / d^2\n",
    "        if j > 1\n",
    "            A[k, kl] -= (reactor2[k].D + reactor2[kl].D) / (2 * d^2)\n",
    "            A[k, k] += reactor2[kl].D / (2 * d^2)\n",
    "        end\n",
    "        if j < N\n",
    "            A[k, kr] -= (reactor2[k].D + reactor2[kr].D) / (2 * d^2)\n",
    "            A[k, k] += reactor2[kr].D / (2 * d^2)\n",
    "        end\n",
    "        if i > 1\n",
    "            A[k, ku] -= (reactor2[k].D + reactor2[ku].D) / (2 * d^2)\n",
    "            A[k, k] += reactor2[ku].D / (2 * d^2)\n",
    "        end\n",
    "        if i < N\n",
    "            A[k, kd] -= (reactor2[k].D + reactor2[kd].D) / (2 * d^2)\n",
    "            A[k, k] += reactor2[kd].D / (2 * d^2)\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    F = reshape(F, N^2)\n",
    "    Sf = [reactor2[i].νΣf for i in eachindex(reactor2)]\n",
    "    λ = 1.0\n",
    "    S = 1 / λ .* Sf .* F\n",
    "    ϵ = 1.0\n",
    "\n",
    "    while ϵ > ϵlim\n",
    "        Fn = A \\ S\n",
    "        λn = sum(Sf .* Fn) / sum(Sf .* F) * λ\n",
    "        Sn = 1 / λn .* Sf .* Fn\n",
    "        ϵ = abs(λn - λ) / λ\n",
    "        F = copy(Fn)\n",
    "        S = copy(Sn)\n",
    "        λ = λn\n",
    "    end\n",
    "    reshape(F, N, N)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7721198d-2fe9-4e6b-813c-1a106a7a215c",
   "metadata": {},
   "source": [
    "Sprawdźmy teraz jak wygląda rozwiązanie dla reaktora bez reflektora (powiększamy Ra kosztem Rr, aby zachować ten sam rozmiar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15dd1b5f-946b-42be-9387-01d9e6d1b93f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Ra = 7\n",
    "Rr = 0\n",
    "Rc = 25\n",
    "N = 2 * (Ra + Rr + Rc + 1)\n",
    "\n",
    "reactor2b = Array{Material, 2}(undef, N, N)\n",
    "for j in 1:N, i in 1:N\n",
    "    r = sqrt((i - N/2)^2 + (j - N/2)^2)\n",
    "    if r > Rr + Rc\n",
    "        reactor2b[i, j] = void\n",
    "    elseif Rr + Rc >= r > Rc\n",
    "        reactor2b[i, j] = moderator\n",
    "    else \n",
    "        reactor2b[i, j] = core\n",
    "    end\n",
    "end\n",
    "\n",
    "fig = Figure(size=(800, 600))\n",
    "Sa = zeros(N, N)\n",
    "for j in 1:N, i in 1:N\n",
    "    Sa[i, j] = reactor2b[i, j].Σa\n",
    "end\n",
    "ax = Axis(fig[1, 1]; xlabel=\"x\", ylabel=\"y\")\n",
    "hm = heatmap!(ax, Sa, colormap=:lighttest)\n",
    "Colorbar(fig[1, 2], hm)\n",
    "fig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ce85eddb-7fb0-4f36-9701-105a73fa1906",
   "metadata": {},
   "outputs": [],
   "source": [
    "F1 = solve2d(reactor2)\n",
    "F2 = solve2d(reactor2b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd28c950-8732-43cd-a5eb-8da9fb69034e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "x0 = round(Int, N/2)\n",
    "fig = Figure(size=(800, 600))\n",
    "ax = Axis(fig[1, 1]; xlabel=\"y\", ylabel=L\"\\Phi\")\n",
    "lines!(ax, F1[x0, :], label=\"refl.\", color=:blue)\n",
    "lines!(ax, F2[x0, :], label=\"bare\", color=:red)\n",
    "vlines!(ax, [x0-Rc, x0+Rc], color=:gray, linestyle=:dash)\n",
    "axislegend(ax)\n",
    "\n",
    "@printf(\"%.3e %.3e\\n\", sum(F1[x0, x0-Rc:x0+Rc]), sum(F2[x0, x0-Rc:x0+Rc]))\n",
    "\n",
    "fig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91c132fa-6220-42e0-b79b-f017854f4c1a",
   "metadata": {},
   "source": [
    "Nasz model reaktora możemy teraz oczywiście dowolnie komplikować (w ramach możliwości obliczeniowych komputera - algorytmy odwracania macierzy mają złożoność $N^2$), wprowadzać różne materiały np. rdzeń o większym i mniejszym wzbogaceniu, pręty kontrolne, siatkę paliwo/moderator itd."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.12",
   "language": "julia",
   "name": "julia-1.12"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}