Modification Kolmolaw

fluidsim/base/output/kolmo_law3d.py

@@ -12,55 +12,59 @@
 import itertools
 
 import numpy as np
-
+import os
 from fluiddyn.util import mpi
 
+import h5py
+import matplotlib.pyplot as plt
+
 from .base import SpecificOutput
 
+from math import floor
+
 """ Conversion from cartesian coordinates system to spherical coordinate system """
 
 
 class OperKolmoLaw:
     def __init__(self, X, Y, Z, params):
-        self.r = np.sqrt(X**2 +Y**2 + Z**2)
+        self.r = np.sqrt(X**2 + Y**2 + Z**2)
         self.rh = np.sqrt(X**2 + Y**2)
         self.rz = Z
-#        self.sin_theta_sin_phi = Y / self.r
-#        self.cos_phi_sin_theta  = X / self.r
-#        self.sin_theta = self.rh / self.r
-#        self.cos_theta = Z / self.r
-#        self.sin_phi = Y  / self.rh
-#        self.cos_phi = X  / self.rh
 
+
 
     def average_azimutal(self, arr):
 
-
         avg_arr = None
         if mpi.nb_proc == 1:
-            avg_arr = np.mean(arr, axis=(0,1))
+            avg_arr = np.mean(arr, axis=(0, 1))
         return avg_arr
 
-        local_sum = np.sum(arr, axis=(0,1))
+        local_sum = np.sum(arr, axis=(0, 1))
         if mpi.rank == 0:
             global_arr = np.zeros(self.nz)  # define array to sum on all proc
 
         for rank in range(mpi.nb_proc):
             if mpi.rank == 0:
                 nz_loc = self.nzs_local[rank]  # define size of array on each proc
+            print("nz_loc " + str(nz_loc))
             if rank == 0 and mpi.rank == 0:
                 sum = local_sum  # start the sum on rank 0
             else:
-                if mpi.rank == 0:  # sum made on rank 0: receive local_array of rank
+                # sum made on rank 0: receive local_array of rank
+                if (
+                    mpi.rank == 0
+                ):  
                     sum = np.empty(nz_loc)
                     mpi.comm.Recv(sum, source=rank, tag=42 * rank)
                 elif mpi.rank == rank:  # send the local array to 0
                     mpi.comm.Send(sum_local, dest=0, tag=42 * rank)
             if mpi.rank == 0:  # construct global sum on 0
                 iz_start = self.izs_start[rank]
+                print("iz_start " + iz_start)
                 global_array[iz_start : iz_start + nz_loc] += sum
         if mpi.rank == 0:
-            avg_arr = global_array / self.nazim
+            avg_arr = global_array / (params.oper.nx * params.oper.ny)
             return avg_arr
 
 
@@ -130,7 +134,7 @@ class KolmoLaw(SpecificOutput):
     """
 
     _tag = "kolmo_law"
-    _name_file = "kolmo_law.h5"
+    _name_file = _tag + ".h5"
 
     @classmethod
     def _complete_params_with_default(cls, params):
@@ -145,31 +149,96 @@ def __init__(self, output):
             period_save_kolmo_law = params.output.periods_save.kolmo_law
         except AttributeError:
             period_save_kolmo_law = 0.0
+        period_save_kolmo_law = 0.1
         if period_save_kolmo_law != 0.0:
             X, Y, Z = output.sim.oper.get_XYZ_loc()
             self.oper_kolmo_law = OperKolmoLaw(X, Y, Z, params)
+
+            self.rhrz = {
+                "rh": self.oper_kolmo_law.rh,
+                "rz": self.oper_kolmo_law.rz
+            }
+            self.rh_max = np.sqrt(params.oper.Lx**2 + params.oper.Ly**2)
+            self.rz_max = params.oper.Lz
+ #          n_sort = params.n_sort
+            self.n_sort = 10
+            n_sort = self.n_sort
+            rh_sort = np.empty([n_sort])
+            rz_sort = np.empty([n_sort])
+            self.drhrz = {
+            "drh": rh_sort,
+            "drz": rz_sort,
+            }
+
+
+            for i in range(n_sort):
+                self.drhrz["drh"][i] = self.rh_max * (i + 1) / n_sort
+                self.drhrz["drz"][i] = self.rz_max * (i + 1) / n_sort
             arrays_1st_time = {
-               "rh": self.oper_kolmo_law.rh,
-               "rz": self.oper_kolmo_law.rz,
+                "rh_sort": self.drhrz["drh"],
+                "rz_sort": self.drhrz["drz"],
             }
+
         else:
             arrays_1st_time = None
+        self.rhrz_sort = arrays_1st_time
 
         super().__init__(
             output,
             period_save=period_save_kolmo_law,
             arrays_1st_time=arrays_1st_time,
         )
 
+    def _init_path_files(self):
+        path_run = self.output.path_run
+        self.path_kolmo_law = path_run + "/_name_file"
+        self.path_file = self.path_kolmo_law
+
+    def _init_files(self, arrays_1st_time=None):
+        dict_J_k_hv = self.compute()
+        if mpi.rank == 0:
+            if not os.path.exists(self.path_kolmo_law):
+                self._create_file_from_dict_arrays(
+                    self.path_kolmo_law, dict_J_k_hv, arrays_1st_time
+                )
+                self.nb_saved_times = 1
+            else:
+                with h5py.File(self.path_kolmo_law, "r") as file:
+                    dset_times = file["times"]
+                    self.nb_saved_times = dset_times.shape[0] + 1               
+                self._add_dict_arrays_to_file(self.path_kolmo_law, dict_J_k_hv)
+
+        self.t_last_save = self.sim.time_stepping.t
+
+    def _online_save(self):
+        """Save the values at one time."""
+        tsim = self.sim.time_stepping.t
+        if tsim - self.t_last_save >= self.period_save:
+            self.t_last_save = tsim
+            dict_J_k_hv = self.compute()
+            if mpi.rank == 0:
+                self._add_dict_arrays_to_file(self.path_kolmo_law, dict_J_k_hv)
+                self.nb_saved_times += 1
+
+    def load(self):
+        path_file=self.path_kolmo_law
+        with h5py.File(path_file, "r") as file:
+            for key in file.keys():
+                J_k_hv[key] = file[key][...]
+        return J_k_hv
+
+    def plot_kolmo_law(self):
+        result = self.load()
+        plottedy = result
 
     def compute(self):
         """compute the values at one time."""
         state = self.sim.state
+        params = self.sim.params
         state_phys = state.state_phys
         state_spect = state.state_spect
-        keys_state_phys=state.keys_state_phys
+        keys_state_phys = state.keys_state_phys
         fft = self.sim.oper.fft
-
         letters = "xyz"
 
         tf_vi = [state_spect.get_var(f"v{letter}_fft") for letter in letters]
@@ -178,11 +247,11 @@ def compute(self):
         tf_K = None
 
         if "b" in keys_state_phys:
-            b= check_strat(state_phys.get_var("b"))
-            tf_b = check_strat(state_spect.get_var("b_fft"))
+            b = state_phys.get_var("b")
+            tf_b = state_spect.get_var("b_fft")
             b2 = b * b
             tf_b2 = fft(b2)
-            tf_bv=[None] * 3
+            tf_bv = [None] * 3
             bv = [item * b for item in vel]
             for index in range(len(bv)):
                 tf_bv[index] = fft(bv[index])
@@ -207,7 +276,10 @@ def compute(self):
         for ind_i in range(3):
             tmp = tf_vi[ind_i] * tf_K.conj()
             if "b" in keys_state_phys:
-                mom = 4 * tf_bv[ind_i].conj() * tf_b + 2 * tf_b2.conj() * tf_vi[ind_i]
+                mom = (
+                    4 * tf_bv[ind_i].conj() * tf_b
+                    + 2 * tf_b2.conj() * tf_vi[ind_i]
+                )
                 mom.real = 0.0
             for ind_j in range(3):
                 tmp += tf_vi[ind_j] * tf_vjvi[ind_i, ind_j].conj()
@@ -216,19 +288,53 @@ def compute(self):
             if "b" in keys_state_phys:
                 J_A_xyz[ind_i] = self.sim.oper.ifft(mom)
 
-        result = {}
-        keys = ["r", "theta", "phi"]
-        for index, key in enumerate(keys):
-            result["J_K_" + key] = self.oper_kolmo_law.average_azimutal(
-                J_K_xyz[index]
+
+        n_sort = self.n_sort
+        J_k_z = np.empty([n_sort, n_sort])
+        J_k_h = np.empty([n_sort, n_sort])
+        J_k_hv_prov = {
+            "J_k_h": J_k_h,
+            "J_k_z": J_k_z,
+        }
+        count = np.ones([n_sort, n_sort])
+        rhrz_sort = self.rhrz_sort
+        J_k_xyz = np.array(J_K_xyz)
+
+        for index, value in np.ndenumerate(J_k_xyz[2]):
+            i = floor(self.rhrz["rh"][index] * n_sort / self.rh_max)
+            j = floor(self.rhrz["rz"][index] * n_sort / self.rz_max)
+            count[i, j] += 1
+            J_k_hv_prov["J_k_z"][i, j] += value
+            J_k_hv_prov["J_k_h"][i, j] += np.sqrt(
+                J_k_xyz[0][index] ** 2 + J_k_xyz[1][index] ** 2
             )
+        for index, value in np.ndenumerate(J_k_hv_prov["J_k_z"]):
+            if count[index] == 0:
+                J_k_hv_prov["J_k_z"][index] = 0.0
+                J_k_hv_prov["J_k_h"][index] = 0.0
+            else:
+                J_k_hv_prov["J_k_z"][index] = value / count[index]
+                J_k_hv_prov["J_k_h"][index] = (
+                    J_k_hv_prov["J_k_h"][index] / count[index]
+                )
+
+        if mpi.nb_proc > 0:
+            collect_J_k_z = mpi.comm.gather(J_k_hv_prov["J_k_z"], root=0)
+            collect_J_k_h = mpi.comm.gather(J_k_hv_prov["J_k_h"], root=0)
+            if mpi.rank == 0:
+                J_k_hv_prov["J_k_z"] = sum(collect_J_k_z) / len(collect_J_k_z)
+                J_k_hv_prov["J_k_h"] = sum(collect_J_k_h) / len(collect_J_k_h)
 
+        result = J_k_hv_prov
+ #       print("rh_sort"+rhrz_sort["rh_sort"])
         return result
 
+
     def check_diff_methods(self):
         first_method = compute(self)
         second_method = compute_alt(self)
         if not np.allclose(first_method, second_method):
             raise RuntimeError(
-                "Both methods are inconsistent: " " ({self.sim.time_stepping.it = })"
+                "Both methods are inconsistent: "
+                " ({self.sim.time_stepping.it = })"
             )

fluidsim/solvers/ns3d/test_solver.py

@@ -63,7 +63,7 @@ def init_params(cls):
         params.oper.coef_dealiasing = 2.0 / 3
         params.nu_4 = 2.0
         params.nu_8 = 2.0
-
+        params.n_sort = 5.0
         params.time_stepping.t_end = 1.5 * params.time_stepping.deltat_max
         params.init_fields.type = "noise"
 
@@ -407,6 +407,8 @@ def customize(result, sim):
         assert df.loc[0, "foo"] > 0
         sim.output.get_mean_values(customize=customize)
 
+        sim.output.kolmo_law.compute()
+
 
 class TestInitInScript(TestSimulBase):
     @classmethod