From d9a7da5afcdfad1d2ad98ab3e0b6a4ab45bf0409 Mon Sep 17 00:00:00 2001
From: paugier <pierre.augier@univ-grenoble-alpes.fr>
Date: Wed, 12 Jun 2024 10:41:12 +0200
Subject: [PATCH] doc/examples/simul_ns3d_forced_isotropic.py and co

---
 doc/examples/simul_ns3d_forced_isotropic.py   | 84 +++++++++++++++++++
 fluidsim/base/forcing/specific.py             |  4 +-
 .../util/scripts/turb_trandom_anisotropic.py  |  2 +-
 lib/fluidsim_core/scripts/restart.py          |  2 +-
 pyproject.toml                                |  9 +-
 5 files changed, 95 insertions(+), 6 deletions(-)
 create mode 100644 doc/examples/simul_ns3d_forced_isotropic.py

diff --git a/doc/examples/simul_ns3d_forced_isotropic.py b/doc/examples/simul_ns3d_forced_isotropic.py
new file mode 100644
index 000000000..8ff4b4cf1
--- /dev/null
+++ b/doc/examples/simul_ns3d_forced_isotropic.py
@@ -0,0 +1,84 @@
+import argparse
+import os
+
+from fluiddyn.util.mpi import printby0
+
+from fluidsim.solvers.ns3d.solver import Simul
+
+parser = argparse.ArgumentParser()
+
+parser.add_argument(
+    "--nx",
+    type=int,
+    default=96,
+    help="Number of grid points in the x direction.",
+)
+parser.add_argument(
+    "--t_end", type=float, default=8.0, help="End time of the simulation"
+)
+parser.add_argument(
+    "--order",
+    type=int,
+    default=4,
+    help="Order of the viscosity (`2` corresponds to standard viscosity)",
+)
+
+args = parser.parse_args()
+
+if "FLUIDSIM_TESTS_EXAMPLES" in os.environ:
+    t_end = 1.0
+    nx = 24
+else:
+    t_end = args.t_end
+    nx = args.nx
+
+params = Simul.create_default_params()
+
+params.output.sub_directory = "examples"
+
+ny = nz = nx
+Lx = 3
+params.oper.nx = nx
+params.oper.ny = ny
+params.oper.nz = nz
+params.oper.Lx = Lx
+params.oper.Ly = Ly = Lx / nx * ny
+params.oper.Lz = Lz = Lx / nx * nz
+
+params.time_stepping.USE_T_END = True
+params.time_stepping.t_end = t_end
+
+order_visco = args.order
+dx = Lx / nx
+epsilon = 1.0
+C = 1.0
+nu = (dx / C) ** ((3 * order_visco - 2) / 3) * epsilon ** (1 / 3)
+setattr(params, f"nu_{order_visco}", nu)
+
+printby0(f"nu_{order_visco} = {nu:.3e}")
+
+params.init_fields.type = "noise"
+params.init_fields.noise.length = 1.0
+params.init_fields.noise.velo_max = 0.1
+
+params.forcing.enable = True
+params.forcing.type = "tcrandom"
+params.forcing.normalized.constant_rate_of = None
+params.forcing.nkmin_forcing = 3
+params.forcing.nkmax_forcing = 4
+# solenoidal field (toroidal + poloidal)
+params.forcing.key_forced = ["vt_fft", "vp_fft"]
+# forcing rate **per key forced**
+params.forcing.forcing_rate = 0.5 * epsilon
+
+params.output.periods_print.print_stdout = 1e-1
+
+params.output.periods_save.phys_fields = 0.5
+params.output.periods_save.spatial_means = 0.1
+params.output.periods_save.spectra = 0.1
+params.output.periods_save.spect_energy_budg = 0.1
+
+params.output.spectra.kzkh_periodicity = 1
+
+sim = Simul(params)
+sim.time_stepping.start()
diff --git a/fluidsim/base/forcing/specific.py b/fluidsim/base/forcing/specific.py
index 8fbaa4118..82f64e457 100644
--- a/fluidsim/base/forcing/specific.py
+++ b/fluidsim/base/forcing/specific.py
@@ -489,11 +489,11 @@ def compute(self):
                 a_fft = self.oper.coarse_seq_from_fft_loc(
                     a_fft, self.shapeK_loc_coarse
                 )
-            except IndexError:
+            except IndexError as error:
                 raise ValueError(
                     f"rank={self.oper.rank}; {self.shapeK_loc_coarse = }; "
                     f"{self.oper.shapeK_loc = }"
-                )
+                ) from error
 
             if mpi.rank == 0:
                 fa_fft = self.forcingc_raw_each_time(a_fft)
diff --git a/fluidsim/util/scripts/turb_trandom_anisotropic.py b/fluidsim/util/scripts/turb_trandom_anisotropic.py
index 734b41800..09aa1ead6 100644
--- a/fluidsim/util/scripts/turb_trandom_anisotropic.py
+++ b/fluidsim/util/scripts/turb_trandom_anisotropic.py
@@ -497,7 +497,7 @@ def main(args=None, **defaults):
 
 # To visualize with fluidsim:
 
-cd {sim.output.path_run}; fluidsim-ipy-load
+fluidsim-ipy-load {sim.output.path_run}
 
 # in IPython:
 
diff --git a/lib/fluidsim_core/scripts/restart.py b/lib/fluidsim_core/scripts/restart.py
index 2d9223822..344260369 100644
--- a/lib/fluidsim_core/scripts/restart.py
+++ b/lib/fluidsim_core/scripts/restart.py
@@ -112,7 +112,7 @@ def restart(self, args=None, **defaults):
         """
         # To visualize with IPython:
 
-        cd {path_run}; fluidsim-ipy-load
+        fluidsim-ipy-load {path_run}
     """
     )
 
diff --git a/pyproject.toml b/pyproject.toml
index 380f7a88a..d4d39d7c0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -163,9 +163,9 @@ dev = [
 ]
 
 [tool.pdm.scripts]
-black = 'black -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
+black = 'black fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
 lint = {shell="pylint -rn --rcfile=pylintrc --jobs=$(nproc) fluidsim --exit-zero"}
-black_check = 'black --check -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
+black_check = 'black --check fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/"'
 validate_code = {composite = ["black_check", "lint"]}
 
 [tool.pdm.options]
@@ -254,3 +254,8 @@ showcontent = true
 directory = "security"
 name = "Security"
 showcontent = true
+
+[tool.black]
+line-length = 82
+target-version = ["py39"]
+extend-exclude = "/(__pythran__|__python__|__numba__|build|doc/_build|\\.ipynb_checkpoints/*)/"