# HG changeset patch
# User paugier <pierre.augier@univ-grenoble-alpes.fr>
# Date 1592412715 -7200
#      Wed Jun 17 18:51:55 2020 +0200
# Node ID ececac5c495b860353ea4a1430c28d0fba9e9051
# Parent  99d73705f72b60f3ee865f39d97a64b74a1c4c45
Profiling and optimizations

diff --git a/doc/examples/time_stepping/3d/analyse_profiles.py b/doc/examples/time_stepping/3d/analyse_profiles.py
--- a/doc/examples/time_stepping/3d/analyse_profiles.py
+++ b/doc/examples/time_stepping/3d/analyse_profiles.py
@@ -30,7 +30,7 @@
 
 
 def get_dataframe(path_dir):
-
+    path_dir = Path(path_dir)
     paths_log = sorted(path_dir.glob("*.log"))
 
     values = []
@@ -39,6 +39,11 @@
 
     df = DataFrame(values, columns=columns)
     df["nx2/3"] = df.nx * df.coef_dealiased / (2 / 3)
-    # norm = df[df.scheme == "RK4"].duration.values[0]
-    # df["speedup"] = norm / df["duration"]
+    norm = df[df.scheme == "RK4"].duration.values[0]
+    df["speedup"] = norm / df["duration"]
     return df
+
+
+df = get_dataframe("tmp_profile")
+
+print(df)
diff --git a/doc/examples/time_stepping/3d/run_profile.py b/doc/examples/time_stepping/3d/run_profile.py
--- a/doc/examples/time_stepping/3d/run_profile.py
+++ b/doc/examples/time_stepping/3d/run_profile.py
@@ -35,10 +35,14 @@
     init_state(sim)
     Path("tmp_profile").mkdir(exist_ok=True)
     base_name_file = f"tmp_profile/{sim.name_run}."
-    with open(base_name_file + "log", "w") as file:
+    name_file_log = base_name_file + "log"
+    with open(name_file_log, "w") as file:
         with redirect_stdout(file):
             run_profile(sim, path_results=base_name_file + "pstats")
 
+    with open(name_file_log) as file:
+        print(file.read())
+
 
 if __name__ == "__main__":
     args = parser.parse_args()
diff --git a/doc/examples/time_stepping/3d/submit_all.py b/doc/examples/time_stepping/3d/submit_all.py
--- a/doc/examples/time_stepping/3d/submit_all.py
+++ b/doc/examples/time_stepping/3d/submit_all.py
@@ -1,18 +1,15 @@
-from which_params import compute, coefficients, nx_bigs
+from which_params import compute, coefficients, nx_bigs, nb_procs
 from util_submit import submit_simul
 
 schemes = [
     "RK4",
     "RK2",
-    "RK2_phaseshift",
     "RK2_phaseshift_exact",
     "RK2_phaseshift_random",
-    ("RK2_phaseshift", 0.2),
-    ("RK2_phaseshift_random", 0.2),
 ]
 
 
-def submit(nx_big):
+def submit(nx_big, nb_proc):
     nb_points = compute(nx_big, verbose=False)
 
     assert nb_points[0] == nx_big
@@ -29,8 +26,10 @@
             else:
                 cfl_coef = None
             print(f"{coef_dealiasing:.3f}, {nx}, {scheme:25s}, {cfl_coef}")
-            submit_simul(coef_dealiasing, nx, scheme, cfl_coef)
+            submit_simul(coef_dealiasing, nx, scheme, cfl_coef, nb_proc)
 
 
-for nx_big in nx_bigs:
-    submit(nx_big)
+for nx_big, nb_proc in zip(nx_bigs, nb_procs):
+    if nx_big > 400:
+        continue
+    submit(nx_big, nb_proc)
diff --git a/doc/examples/time_stepping/3d/submit_profiles.py b/doc/examples/time_stepping/3d/submit_profiles.py
--- a/doc/examples/time_stepping/3d/submit_profiles.py
+++ b/doc/examples/time_stepping/3d/submit_profiles.py
@@ -20,7 +20,7 @@
     coefficients = [0.9, 1.0]
 
     for nx, coef_dealiasing in zip(nxs[1:], coefficients):
-        for scheme in ["RK2_phaseshift", "RK2_phaseshift_random"]:
+        for scheme in ["RK2_phaseshift_random", "RK2_phaseshift_exact"]:
             submit_profile(coef_dealiasing, nx, scheme, t_end=t_end)
 
 
diff --git a/doc/examples/time_stepping/3d/util_submit.py b/doc/examples/time_stepping/3d/util_submit.py
--- a/doc/examples/time_stepping/3d/util_submit.py
+++ b/doc/examples/time_stepping/3d/util_submit.py
@@ -11,17 +11,18 @@
     "export FLUIDSIM_PATH=/fsnet/project/watu/2020/20MILESTONE",
 ]
 
-nb_nodes = 1
 
-
-def submit_simul(coef_dealiasing, nx, type_time_scheme, cfl_coef=None):
+def submit_simul(
+    coef_dealiasing, nx, type_time_scheme, cfl_coef=None, nb_proc=None
+):
 
-    if nx < 480:
-        nb_cores_per_node = 10
-    else:
-        nb_cores_per_node = 20
+    if nb_proc is None:
+        if nx < 480:
+            nb_proc = 10
+        else:
+            nb_proc = 20
 
-    nb_mpi_processes = nb_nodes * nb_cores_per_node
+    nb_mpi_processes = nb_proc
 
     command = (
         f"run_simul.py -cd {coef_dealiasing} -nx {nx} --type_time_scheme {type_time_scheme} "
@@ -45,8 +46,8 @@
         cluster.submit_script(
             command,
             name_run=name_run,
-            nb_nodes=nb_nodes,
-            nb_cores_per_node=nb_cores_per_node,
+            nb_nodes=1,
+            nb_cores_per_node=nb_proc,
             nb_mpi_processes=nb_mpi_processes,
             omp_num_threads=1,
             idempotent=True,
@@ -64,8 +65,6 @@
     else:
         nb_cores_per_node = 20
 
-    nb_mpi_processes = nb_nodes * nb_cores_per_node
-
     command = (
         f"run_profile.py -cd {coef_dealiasing} -nx {nx} "
         f"--type_time_scheme {type_time_scheme} --t_end {t_end}"
@@ -88,7 +87,7 @@
         cluster.submit_script(
             command,
             name_run=name_run,
-            nb_nodes=nb_nodes,
+            nb_nodes=1,
             nb_cores_per_node=nb_cores_per_node,
             nb_mpi_processes=1,
             omp_num_threads=1,
diff --git a/doc/examples/time_stepping/3d/which_params.py b/doc/examples/time_stepping/3d/which_params.py
--- a/doc/examples/time_stepping/3d/which_params.py
+++ b/doc/examples/time_stepping/3d/which_params.py
@@ -1,8 +1,26 @@
 import numpy as np
 
+try:
+    from sympy import primefactors
+except ImportError:
+    pass
+
+L = 1.0
+V0 = 1.0
+Re = 1600
+Lx = 2 * np.pi * L
+deltak = 2 * np.pi / Lx
+nu = V0 * L / Re
+
+# A result from a simulation
+epsilon = 0.012
+# epsilon = 0.008
+eta = (nu ** 3 / epsilon) ** (1 / 4)
+
 coefficients = np.array([2 / 3, 0.90, 1.0])
 
-nx_bigs = [800, 560, 420]
+nx_bigs = [800, 400, 256, 128]
+nb_procs = [20, 10, 4, 2]
 
 
 def compute(nx, nb_processes=20, verbose=True):
@@ -12,23 +30,36 @@
     nb_points[0] = nx
 
     tmp = nx * coefficients[0] / coefficients / (2 * nb_processes) + 1e-10
+    round_tmp = np.round(tmp)
+    if 13 in round_tmp:
+        index = np.where(round_tmp == 13)[0][0]
+        round_tmp[index] = 14
 
-    nb_points[1:] = 2 * nb_processes * np.round(tmp[1:])
+    # if 11 in round_tmp:
+    #     index = np.where(round_tmp == 11)[0][0]
+    #     round_tmp[index] = 10
+
+    nb_points[1:] = 2 * nb_processes * round_tmp[1:]
 
     if verbose:
+        factors = []
+        for number in round_tmp:
+            factors.append(primefactors(int(number)))
         print(f"\n{nb_processes = }")
-        print("coef dealisasing", coefficients)
-        # print(tmp)
-        # print(np.round(tmp))
         print("nx =", nb_points)
         kmax = coefficients * deltak * nb_points / 2
-        print(f"{kmax[0] = :.2f}")
+        print(f"{kmax[0]*eta = :.2f}")
+        print(tmp)
+        print(round_tmp)
+        print(f"{factors = }")
         print(np.round(100 * (kmax - kmax[0]) / kmax[0], decimals=2), "%")
 
     return nb_points
 
 
 if __name__ == "__main__":
-    compute(nx_bigs[0])
-    compute(nx_bigs[1])
-    compute(nx_bigs[2], 10)
+
+    print("coef dealisasing", coefficients)
+
+    for nx, nb_proc in zip(nx_bigs, nb_procs):
+        compute(nx, nb_proc)
diff --git a/fluidsim/base/time_stepping/pseudo_spect.py b/fluidsim/base/time_stepping/pseudo_spect.py
--- a/fluidsim/base/time_stepping/pseudo_spect.py
+++ b/fluidsim/base/time_stepping/pseudo_spect.py
@@ -51,6 +51,7 @@
     state_spect: A, dt: float, tendencies: A, diss: ArrayDiss, output: A
 ):
     output[:] = (state_spect + dt * tendencies) * diss
+    return output
 
 
 @boost
@@ -65,6 +66,26 @@
     state_spect[:] = state_spect * diss + dt * diss2 * tendencies
 
 
+@boost
+def mean_with_phaseshift(
+    tendencies_0: A, tendencies_1_shift: A, phase_shift: Am1, output: A
+):
+    output[:] = 0.5 * (tendencies_0 + tendencies_1_shift / phase_shift)
+    return output
+
+
+@boost
+def mul(phase_shift: Am1, state_spect: A, output: A):
+    output[:] = phase_shift * state_spect
+    return output
+
+
+@boost
+def div_inplace(arr: A, phase_shift: Am1):
+    arr /= phase_shift
+    return arr
+
+
 class ExactLinearCoefs:
     """Handle the computation of the exact coefficient for the RK4."""
 
@@ -604,28 +625,30 @@
         compute_tendencies = self.sim.tendencies_nonlin
         state_spect = self.sim.state.state_spect
 
-        tendencies_0 = (
-            compute_tendencies(phase_shift_alpha * state_spect)
-            / phase_shift_alpha
+        tmp0 = np.empty_like(state_spect)
+        state_spect_shift = mul(phase_shift_alpha, state_spect, output=tmp0)
+        tendencies_0_shift = compute_tendencies(
+            state_spect_shift, old=state_spect_shift
+        )
+
+        tendencies_0 = div_inplace(tendencies_0_shift, phase_shift_alpha)
+
+        state_spect_1 = step_Euler(
+            state_spect, dt, tendencies_0, diss, output=self._state_spect_tmp
         )
 
-        state_spect_1 = self._state_spect_tmp
-        step_Euler(state_spect, dt, tendencies_0, diss, output=state_spect_1)
-
-        tendencies_1_shift = compute_tendencies(phase_shift_beta * state_spect_1)
+        tmp1 = np.empty_like(state_spect)
+        state_spect_1_shift = mul(phase_shift_beta, state_spect_1, output=tmp1)
+        tendencies_1_shift = compute_tendencies(
+            state_spect_1_shift, old=state_spect_1_shift
+        )
 
-        tendencies_d = self._state_spect_tmp
-        if ts.is_transpiled:
-            ts.use_block("rk2_random")
-        else:
-            # transonic block (
-            #     A tendencies_d, tendencies_0, tendencies_1_shift;
-            #     Am1 phase_shift_beta;
-            # )
-            tendencies_d[:] = 0.5 * (
-                tendencies_0 + tendencies_1_shift / phase_shift_beta
-            )
-
+        tendencies_d = mean_with_phaseshift(
+            tendencies_0,
+            tendencies_1_shift,
+            phase_shift_beta,
+            output=self._state_spect_tmp,
+        )
         step_like_RK2(state_spect, dt, tendencies_d, diss, diss2)
 
     def _time_step_RK2_phaseshift_exact(self):
@@ -669,25 +692,49 @@
         """
         dt = self.deltat
         diss, diss2 = self.exact_linear_coefs.get_updated_coefs()
-
         phase_shift = self._get_phase_shift()
         compute_tendencies = self.sim.tendencies_nonlin
         state_spect = self.sim.state.state_spect
 
-        tendencies_d0 = 0.5 * (
-            compute_tendencies(state_spect)
-            + compute_tendencies(phase_shift * state_spect) / phase_shift
+        tmp0 = np.empty_like(state_spect)
+        tmp1 = np.empty_like(state_spect)
+        tmp2 = np.empty_like(state_spect)
+        tmp3 = np.empty_like(state_spect)
+
+        tendencies_0 = compute_tendencies(state_spect, old=tmp0)
+        state_spect_shift = mul(phase_shift, state_spect, output=tmp1)
+        tendencies_0_shift = compute_tendencies(
+            state_spect_shift, old=state_spect_shift
+        )
+        tendencies_d0 = mean_with_phaseshift(
+            tendencies_0, tendencies_0_shift, phase_shift, output=tmp2
+        )
+
+        state_spect_1 = step_Euler(
+            state_spect, dt, tendencies_d0, diss, output=self._state_spect_tmp
         )
 
-        state_spect_1 = self._state_spect_tmp
-        step_Euler(state_spect, dt, tendencies_d0, diss, output=state_spect_1)
-
-        tendencies_d1 = 0.5 * (
-            compute_tendencies(state_spect_1)
-            + compute_tendencies(phase_shift * state_spect_1) / phase_shift
+        tendencies_1 = compute_tendencies(state_spect_1, old=tmp0)
+        state_spect_shift = mul(phase_shift, state_spect_1, output=tmp1)
+        tendencies_1_shift = compute_tendencies(
+            state_spect_shift, old=state_spect_shift
         )
 
-        tendencies_d = (tendencies_d0 + tendencies_d1) / 2
+        tendencies_d = tmp3
+        if ts.is_transpiled:
+            ts.use_block("rk2_exact")
+        else:
+            # based on approximation 1
+            # transonic block (
+            #     A tendencies_d, tendencies_d0, tendencies_1, tendencies_1_shift;
+            #     Am1 phase_shift
+            # )
+
+            tendencies_d[:] = 0.5 * (
+                tendencies_d0
+                + 0.5 * (tendencies_1 + tendencies_1_shift / phase_shift)
+            )
+
         step_like_RK2(state_spect, dt, tendencies_d, diss, diss2)
 
     def _time_step_RK4(self):
@@ -791,20 +838,14 @@
         state_spect = self.sim.state.state_spect
 
         tendencies_0 = compute_tendencies()
-        state_spect_tmp = self._state_spect_tmp
         state_spect_tmp1 = self._state_spect_tmp1
-        state_spect_12_approx1 = state_spect_tmp1
 
         # rk4_step0
-        step_Euler(
-            state_spect, dt / 6, tendencies_0, diss, output=state_spect_tmp
+        state_spect_tmp = step_Euler(
+            state_spect, dt / 6, tendencies_0, diss, output=self._state_spect_tmp
         )
-        step_Euler(
-            state_spect,
-            dt / 2,
-            tendencies_0,
-            diss2,
-            output=state_spect_12_approx1,
+        state_spect_12_approx1 = step_Euler(
+            state_spect, dt / 2, tendencies_0, diss2, output=state_spect_tmp1,
         )
 
         tendencies_1 = compute_tendencies(
diff --git a/fluidsim/operators/operators3d.py b/fluidsim/operators/operators3d.py
--- a/fluidsim/operators/operators3d.py
+++ b/fluidsim/operators/operators3d.py
@@ -80,6 +80,11 @@
     ff_fft[np.nonzero(where_dealiased)] = 0.0
 
 
+@boost
+def compute_energy_from_1field(arr: Ac):
+    return 0.5 * np.abs(arr) ** 2
+
+
 if not ts.is_transpiling and not ts.is_compiled and not _is_testing:
     # for example if Pythran is not available
     dealiasing_variable = dealiasing_variable_numpy
diff --git a/fluidsim/solvers/nl1d/test_solver.py b/fluidsim/solvers/nl1d/test_solver.py
--- a/fluidsim/solvers/nl1d/test_solver.py
+++ b/fluidsim/solvers/nl1d/test_solver.py
@@ -94,6 +94,10 @@
         s_exact2_fft = oper2.fft(s_exact2)
         cls.s_exact2_fft = s_exact2_fft[oper2.kx <= sim.oper.kx.max()]
 
+    @classmethod
+    def tearDownClass(cls):
+        cls.sim.time_stepping.finalize_main_loop()
+
     def _test_type_time_scheme(self, type_time_scheme, coef_dealiasing=0.66):
         sim = self.sim
         params = sim.params
diff --git a/fluidsim/solvers/ns3d/output/__init__.py b/fluidsim/solvers/ns3d/output/__init__.py
--- a/fluidsim/solvers/ns3d/output/__init__.py
+++ b/fluidsim/solvers/ns3d/output/__init__.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 from fluidsim.base.output import OutputBasePseudoSpectral
+from fluidsim.operators.operators3d import compute_energy_from_1field
 
 
 class Output(OutputBasePseudoSpectral):
@@ -74,9 +75,9 @@
         vy_fft = self.sim.state.state_spect.get_var("vy_fft")
         vz_fft = self.sim.state.state_spect.get_var("vz_fft")
         return (
-            0.5 * np.abs(vx_fft) ** 2,
-            0.5 * np.abs(vy_fft) ** 2,
-            0.5 * np.abs(vz_fft) ** 2,
+            compute_energy_from_1field(vx_fft),
+            compute_energy_from_1field(vy_fft),
+            compute_energy_from_1field(vz_fft),
         )
 
     def compute_energy_fft(self):