Merge pull request #1 from envfluids/init_as_function

Init as function

py2d/Py2D_solver.py

@@ -28,6 +28,7 @@
 from py2d.aposteriori_analysis import eddyTurnoverTime_2DFHIT
 from py2d.SGSModel import *
 # from py2d.uv2tau_CNN import *
+from py2d.initialize import *
 
 # Enable x64 Precision for Jax
 jax.config.update('jax_enable_x64', True)
@@ -220,25 +221,7 @@ def Py2D_solver(Re, fkx, fky, alpha, beta, NX, SGSModel_string, eddyViscosityCoe
     if readTrue:
 
         # -------------- Initialization using pertubration --------------
-        kp = 10.0
-        A  = 4*np.power(kp,(-5))/(3*np.pi)
-        absK = np.sqrt(Kx*Kx+Ky*Ky)
-        Ek = A*np.power(absK,4)*np.exp(-np.power(absK/kp,2))
-        coef1 = np.random.uniform(0,1,[NX//2+1,NX//2+1])*np.pi*2
-        coef2 = np.random.uniform(0,1,[NX//2+1,NX//2+1])*np.pi*2
-
-        perturb = np.zeros([NX,NX])
-        perturb[0:NX//2+1, 0:NX//2+1] = coef1[0:NX//2+1, 0:NX//2+1]+coef2[0:NX//2+1, 0:NX//2+1]
-        perturb[NX//2+1:, 0:NX//2+1] = coef2[NX//2-1:0:-1, 0:NX//2+1] - coef1[NX//2-1:0:-1, 0:NX//2+1]
-        perturb[0:NX//2+1, NX//2+1:] = coef1[0:NX//2+1, NX//2-1:0:-1] - coef2[0:NX//2+1, NX//2-1:0:-1]
-        perturb[NX//2+1:, NX//2+1:] = -(coef1[NX//2-1:0:-1, NX//2-1:0:-1] + coef2[NX//2-1:0:-1, NX//2-1:0:-1])
-        perturb = np.exp(1j*perturb)
-
-        w1_hat = np.sqrt(absK/np.pi*Ek)*perturb*np.power(NX,2)
-
-        psi_hat = -w1_hat*invKsq
-        psiPrevious_hat = psi_hat.astype(np.complex128)
-        psiCurrent_hat = psi_hat.astype(np.complex128)
+        w1_hat, psi_hat, psiPrevious_hat, psiCurrent_hat = initialize_perturbation(NX, Kx, Ky):
         time = 0.0
 
     else:
@@ -481,4 +464,4 @@ def Py2D_solver(Re, fkx, fky, alpha, beta, NX, SGSModel_string, eddyViscosityCoe
     # Print elapsed time
     print('Total Iteration: ', it+1)
     endTime = timer()
-    print('Total Time Taken: ', endTime-startTime)
+    print('Total Time Taken: ', endTime-startTime)

py2d/initialize.py

@@ -0,0 +1,31 @@
+import numpy as np
+import numpy as nnp
+import jax.numpy as np
+
+def initialize_perturbation(NX, Kx, Ky):
+    # -------------- Initialization using pertubration --------------
+    kp = 10.0
+    A  = 4*np.power(kp,(-5))/(3*np.pi)
+    absK = np.sqrt(Kx*Kx+Ky*Ky)
+    Ek = A*np.power(absK,4)*np.exp(-np.power(absK/kp,2))
+    coef1 = np.random.uniform(0,1,[NX//2+1,NX//2+1])*np.pi*2
+    coef2 = np.random.uniform(0,1,[NX//2+1,NX//2+1])*np.pi*2
+
+    perturb = np.zeros([NX,NX])
+    perturb[0:NX//2+1, 0:NX//2+1] = coef1[0:NX//2+1, 0:NX//2+1]+coef2[0:NX//2+1, 0:NX//2+1]
+    perturb[NX//2+1:, 0:NX//2+1] = coef2[NX//2-1:0:-1, 0:NX//2+1] - coef1[NX//2-1:0:-1, 0:NX//2+1]
+    perturb[0:NX//2+1, NX//2+1:] = coef1[0:NX//2+1, NX//2-1:0:-1] - coef2[0:NX//2+1, NX//2-1:0:-1]
+    perturb[NX//2+1:, NX//2+1:] = -(coef1[NX//2-1:0:-1, NX//2-1:0:-1] + coef2[NX//2-1:0:-1, NX//2-1:0:-1])
+    perturb = np.exp(1j*perturb)
+
+    w1_hat = np.sqrt(absK/np.pi*Ek)*perturb*np.power(NX,2)
+
+    psi_hat = -w1_hat*invKsq
+    psiPrevious_hat = psi_hat.astype(np.complex128)
+    psiCurrent_hat = psi_hat.astype(np.complex128)
+
+    return w1_hat, psi_hat, psiPrevious_hat, psiCurrent_hat
+
+
+#def initialize_resume(NX, Kx, Ky):
+