Merge branch 'topic/default/plot-forcing-region' into 'branch/default'

fluidsim/base/forcing/specific.py: add plot_forcing_region function to tc_random

See merge request fluiddyn/fluidsim!455

Author: paugier

fluidsim/base/forcing/anisotropic.py

@@ -264,11 +264,10 @@ def plot_forcing_region(self):
         ax.set_aspect("equal")
 
         title = (
-            pforcing.type
-            + "; "
-            + rf"$nk_{{min}} = {pforcing.nkmin_forcing} \delta k_v$; "
-            + rf"$nk_{{max}} = {pforcing.nkmax_forcing} \delta k_v$; "
-            + "\n"
+            f"{pforcing.type}; "
+            rf"$nk_{{min}} = {pforcing.nkmin_forcing} \delta k_v$; "
+            rf"$nk_{{max}} = {pforcing.nkmax_forcing} \delta k_v$;"
+            "\n"
             + r"$\theta_f = {:.0f}^\circ$; ".format(degrees(self.angle))
             + rf"Forced modes = {self.nb_forced_modes}"
         )
@@ -458,23 +457,10 @@ def plot_forcing_region(self):
             ax.text(loc_label_x, loc_label_y, r"$\delta \theta_f$")
 
         # Plot forced modes in red
-        indices_forcing = np.argwhere(self.COND_NO_F == False)
+        indices_forcing = np.argwhere(~self.COND_NO_F)
         for i, index in enumerate(indices_forcing):
-            if ndim == 2:
-                ax.plot(
-                    Kh[0, index[1]],
-                    Kv[index[0], 0],
-                    "ro",
-                    label="Forced mode" if i == 0 else "",
-                )
-            else:
-                ax.plot(
-                    Kh[0, index[1], index[2]],
-                    Kv[index[0], 0, 0],
-                    "ro",
-                    label="Forced mode" if i == 0 else "",
-                )
-
+            label = "Forced mode" if i == 0 else None
+            ax.plot(Kh[*index], Kv[*index], "ro", label=label)
         ax.grid(linestyle="--", alpha=0.4)
         ax.legend()
 

fluidsim/base/forcing/specific.py

@@ -48,6 +48,8 @@
 from pathlib import Path
 
 import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
 
 from fluiddyn.util import mpi
 from fluiddyn.calcul.easypyfft import fftw_grid_size
@@ -866,6 +868,61 @@ def __init__(self, sim):
             else:
                 self.period_change_f0f1 = time_correlation
 
+    def plot_forcing_region(self):
+        """Plots the forcing region"""
+        pforcing = self.params.forcing
+        kf_max = self.kmax_forcing
+
+        try:
+            self.params.oper.nz
+        except AttributeError:
+            ndim = 2
+        else:
+            ndim = 3
+
+        if ndim == 2:
+            Kh = self.oper_coarse.KX
+            Kv = self.oper_coarse.KY
+            deltakv = self.oper.deltaky
+        else:
+            Kh = np.sqrt(self.oper_coarse.Kx**2 + self.oper_coarse.Ky**2)
+            Kv = self.oper_coarse.Kz
+            deltakv = self.oper.deltakz
+
+        fig, ax = plt.subplots()
+        ax.set_aspect("equal")
+
+        title = (
+            f"{pforcing.type}; "
+            rf"$nk_{{min}} = {pforcing.nkmin_forcing} \delta k_v$; "
+            rf"$nk_{{max}} = {pforcing.nkmax_forcing} \delta k_v$;"
+            f"\nForced modes = {self.nb_forced_modes}"
+        )
+
+        ax.set_title(title)
+        ax.set_xlabel(r"$k_h$")
+        ax.set_ylabel(r"$k_z$")
+
+        # Parameters figure
+
+        # Set limits to 120% of the kf_max
+        factor = 1.2
+        ax.set_xlim([0.0, factor * kf_max])
+        ax.set_ylim([0.0, factor * kf_max])
+
+        xticks = np.arange(0.0, factor * kf_max, deltakv)
+        yticks = np.arange(0.0, factor * kf_max, deltakv)
+        ax.set_xticks(xticks)
+        ax.set_yticks(yticks)
+
+        # Plot forced modes in red
+        indices_forcing = np.argwhere(~self.COND_NO_F)
+        for i, index in enumerate(indices_forcing):
+            label = "Forced mode" if i == 0 else ""
+            ax.plot(Kh[*index], Kv[*index], "ro", label=label)
+        ax.grid(linestyle="--", alpha=0.4)
+        ax.legend()
+
     def forcingc_raw_each_time(self, a_fft):
         """Return a coarse forcing as a linear combination of 2 random arrays
 

fluidsim/solvers/ns2d/strat/test_solver.py

@@ -114,6 +114,9 @@ def test_(self):
             P_tot = means["PK_tot"] + means["PA_tot"]
             assert np.allclose(P_tot, self.sim.params.forcing.forcing_rate)
 
+        if mpi.nb_proc == 1:
+            self.sim.forcing.forcing_maker.plot_forcing_region()
+
 
 class TestForcingConstantRateEnergyAP(TestForcingConstantRateEnergy):
     @classmethod

fluidsim/solvers/ns3d/test_solver.py

@@ -135,6 +135,9 @@ def test_forcing(self):
         assert kx_max >= kf_max
         assert ky_max >= kf_max
 
+        if mpi.nb_proc == 1:
+            sim.forcing.forcing_maker.plot_forcing_region()
+
 
 class TestOutput(TestSimulBase):
     @classmethod

pyproject.toml

@@ -269,5 +269,5 @@ showcontent = true
 
 [tool.ruff]
 line-length = 82
-target-version = "py39"
+target-version = "py311"
 exclude = [".venv", "build", ".*", "**/__py*__", "**/__numba__", "doc/_build", "**/.ipynb_checkpoints"]