<jupyter_start><jupyter_text><jupyter_code>import numpy as np
import numba as nb
import typing as tp
import matplotlib.pyplot as plt
from matplotlib import animation 

@nb.njit
def setBounds(S: np.array, bType: int):
  Ni = S.shape[0]; Nj = S.shape[1] 
  if bType is 1:
    S[:, 0] = -S[:, 1];    S[:, Nj-1] = -S[:, Nj-2];    S[0, :] = S[1, :];      S[Ni-1, :] = S[Ni-2, :]; 
  elif bType is 2:
    S[:, 0] = S[:, 1];     S[:, Nj-1] = S[:, Nj-2];     S[0, :] = -S[1, :];     S[Ni-1, :] = -S[Ni-2, :]; 
  else:
    S[:, 0] = S[:, 1];     S[:, Nj-1] = S[:, Nj-2];     S[0, :] = S[1, :];      S[Ni-1, :] = S[Ni-2, :]; 

  S[0, 0] = 0.5 * (S[1, 0] + S[0, 1]); S[0, Nj-1] = 0.5 * (S[1, Nj-1] + S[0, Nj-2])
  S[Ni-1, Nj-1] = 0.5 * (S[Ni-2, Nj-1] + S[Ni-1, Nj-2]); S[Ni-1, 0] = 0.5 * (S[Ni-1, 1] + S[Ni-2, 0])

@nb.njit
def Advect(S1: np.array, S0: np.array, u: np.array, v: np.array, boundType: int, dh: float, dt: float):
  Ni = S0.shape[0]; Nj = S0.shape[1]
  dtp = dt/dh
  for i in nb.prange(1, Ni - 1):
    for j in np.arange(1, Nj - 1):
       jp = j - dtp * u[i, j]; ip = i + dtp * v[i, j] # it makes sense
       if jp < 0.5 : jp = 0.5
       if ip < 0.5 : ip = 0.5
       if jp > Nj - 1.5 : jp = Nj - 1.5
       if ip > Ni - 1.5 : ip = Ni - 1.5;
       iTop = int(ip); jLeft = int(jp); iBottom = int(ip+1); jRight = int(jp+1)
       interpY = (ip-iTop)/(iBottom-iTop); interpX = (jp-jLeft)/(jRight-jLeft)
       SL = S0[iTop, jLeft]  * (1 - interpY) + S0[iBottom, jLeft]  *  interpY
       SR = S0[iTop, jRight] * (1 - interpY) + S0[iBottom, jRight] *  interpY
       S1[i, j] =         SL * (1 - interpX) +                  SR *  interpX
  setBounds(S1, boundType)

@nb.njit
def Project(u: np.array, v: np.array, div: np.array, p: np.array, dh: float):
  Ni = u.shape[0]; Nj = u.shape[1]; dhp = 1/dh
  for i in nb.prange(1, Ni - 1):
    for j in np.arange(1, Nj - 1):
      div[i, j] = u[i, j+1] - u[i, j-1] + v[i-1, j] - v[i+1, j]
  setBounds(p, 0)
  for k in np.arange(20):
    for i in np.arange(1, Ni - 1):
      for j in np.arange(1, Nj - 1):
        p[i, j] = (-0.5 * dh * div[i, j] + p[i+1, j] + p[i-1, j] + p[i, j+1] + p[i, j-1])/4
    setBounds(p, 0)
  for i in np.arange(1, Ni - 1):
      for j in np.arange(1, Nj - 1):
        u[i, j] -= 0.5 * (p[i, j+1] - p[i, j-1]) * dhp
        v[i, j] -= 0.5 * (p[i-1, j] - p[i+1, j]) * dhp
  setBounds(u, 1); setBounds(v, 2)

@nb.njit
def noSource(S: np.array, dt: np.array, t: np.array):
  pass

class Fluid:
  def __init__(self, Ni: int, Nj: int, dh: float, u0: np.array, v0: np.array, p0: np.array, s0: np.array, visc = 1, viscS = 1, uSource = noSource, vSource = noSource, sSource = noSource):
    self._u1 = u0; self._u0 = np.zeros((Ni+2, Nj+2)); self._uSource = uSource
    self._v1 = v0; self._v0 = np.zeros((Ni+2, Nj+2)); self._vSource = vSource

    self._p = p0
    self._div = np.zeros((Ni+2, Nj+2))
    self._s1 = s0; self._s0 = np.zeros((Ni+2, Nj+2)); self._sSource = sSource

    self._dt = 1
    self._t = 0
    self._dh = dh
    self._visc = 1; self._viscS = viscS
    self._dens = 1

  def vStep(self):
    self._uSource(self._u1, self._dt, self._t); self._vSource(self._v1, self._dt, self._t) 
    setBounds(self._u1, 1); setBounds(self._v1, 2);
    self._u0 = self._u1; self._v0 = self._v1 #
    Advect(self._u1, self._u0, self._u0, self._v0, 1, self._dh, self._dt)
    Advect(self._v1, self._v0, self._u0, self._v0, 2, self._dh, self._dt)
    Project(self._u1, self._v1, self._div, self._p, self._dh)

  def sStep(self):
    self._sSource(self._s1, self._dt, self._t)
    setBounds(self._s1, 0)
    self._s0 = self._s1
    Advect(self._s1, self._s0, self._u1, self._v1, 0, self._dh, self._dt)
  
  def step(self, dt):
    self._dt = dt
    self.vStep()
    self.sStep()
    self._t += self._dt<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.2160168598524763, 'fraction_numerical': 0.051896733403582716, 'mean_word_length': 2.3161572052401747, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '216740', 'license': 'MIT', 'n_tokens_mistral': 1939, 'n_tokens_neox': 1854, 'n_words': 452, 'path': 'FluidWorks.ipynb', 'repo_name': 'papipopapu/NeuralNetwork', 'size': '3796'}