Fix Kolmogorov Stepper (#3)

* Fix scaling of injection

* Fix docstring formatting

* Fix docstring formatting

* Add more notes

* Fix docstring

* Add notes to Kolmogorov stepper

exponax/nonlin_fun/_vorticity_convection.py

@@ -120,25 +120,28 @@ def __init__(
             f = -k (2π/L) γ cos(k (2π/L) x₁)
         ```
 
-        i.e., energy of intensity `γ` is injected at wavenumber `k`.
+        i.e., energy of intensity `γ` is injected at wavenumber `k`. Note that
+        the forcing is on the **vorticity**. As such, we get the prefactor `k
+        (2π/L)` and the `sin(...)` turns into a `-cos(...)` (minus sign because
+        the vorticity is derived via the curl).
 
         **Arguments:**
             - `num_spatial_dims`: The number of spatial dimensions `d`.
             - `num_points`: The number of points `N` used to discretize the
-                domain. This **includes** the left boundary point and **excludes**
-                the right boundary point. In higher dimensions; the number of
-                points in each dimension is the same.
-            - `convection_scale`: The scale `b` of the convection term. Defaults to
-                `1.0`.
+                domain. This **includes** the left boundary point and
+                **excludes** the right boundary point. In higher dimensions; the
+                number of points in each dimension is the same.
+            - `convection_scale`: The scale `b` of the convection term. Defaults
+                to `1.0`.
             - `injection_mode`: The wavenumber `k` at which energy is injected.
                 Defaults to `4`.
-            - `injection_scale`: The intensity `γ` of the injection term. Defaults
-                to `1.0`.
-            - `derivative_operator`: A complex array of shape `(d, ..., N//2+1)` that
-                represents the derivative operator in Fourier space.
-            - `dealiasing_fraction`: The fraction of the highest resolved modes that
-                are not aliased. Defaults to `2/3` which corresponds to Orszag's 2/3
-                rule.
+            - `injection_scale`: The intensity `γ` of the injection term.
+                Defaults to `1.0`.
+            - `derivative_operator`: A complex array of shape `(d, ..., N//2+1)`
+                that represents the derivative operator in Fourier space.
+            - `dealiasing_fraction`: The fraction of the highest resolved modes
+                that are not aliased. Defaults to `2/3` which corresponds to
+                Orszag's 2/3 rule.
         """
         super().__init__(
             num_spatial_dims,
@@ -148,13 +151,15 @@ def __init__(
             dealiasing_fraction=dealiasing_fraction,
         )
 
-        # TODO: shouldn't this be scaled differently sine we are in the
-        # streamfunction-vorticity formulation?
         wavenumbers = build_wavenumbers(num_spatial_dims, num_points)
         injection_mask = (wavenumbers[0] == 0) & (wavenumbers[1] == injection_mode)
         self.injection = jnp.where(
             injection_mask,
-            injection_scale * build_scaling_array(num_spatial_dims, num_points),
+            # Need to additional scale the `injection_scale` with the
+            # `injection_mode`, because we apply the forcing on the vorticity.
+            -injection_mode
+            * injection_scale
+            * build_scaling_array(num_spatial_dims, num_points),
             0.0,
         )
 

exponax/stepper/_navier_stokes.py

@@ -201,6 +201,35 @@ def __init__(
         A negative drag coefficient `λ` is needed to remove some of the energy
         piling up in low modes.
 
+        According to
+
+            Chandler, G.J. and Kerswell, R.R. (2013) ‘Invariant recurrent
+            solutions embedded in a turbulent two-dimensional Kolmogorov flow’,
+            Journal of Fluid Mechanics, 722, pp. 554–595.
+            doi:10.1017/jfm.2013.122.
+
+        equation (2.5), the Reynolds number of the Kolmogorov flow is given by
+
+            Re = √ζ / ν √(L / (2π))³
+
+        with `ζ` being the scaling of the Kolmogorov forcing, i.e., the
+        `injection_scale`. Hence, in the case of `L = 2π`, `ζ = 1`, the Reynolds
+        number is `Re = 1 / ν`. If one uses the default value of `ν = 0.001`,
+        the Reynolds number is `Re = 1000` which also corresponds to the main
+        experiments in
+
+            Kochkov, D., Smith, J.A., Alieva, A., Wang, Q., Brenner, M.P. and
+            Hoyer, S., 2021. Machine learning–accelerated computational fluid
+            dynamics. Proceedings of the National Academy of Sciences, 118(21),
+            p.e2101784118.
+
+        together with `injection_mode = 4`. Note that they required a resolution
+        of `num_points = 2048` (=> 2048^2 = 4.2M degrees of freedom in 2d) to
+        fully resolve all scales at that Reynolds number. Using `Re = 0.01`
+        which corresponds to `ν = 0.01` can be a good starting for
+        `num_points=128`.
+
+
         **Arguments:**
             - `num_spatial_dims`: The number of spatial dimensions `d`.
             - `domain_extent`: The size of the domain `L`; in higher dimensions
@@ -218,8 +247,8 @@ def __init__(
                 convection term. Default is `1.0`.
             - `drag`: The drag coefficient `λ`. Default is `-0.1`.
             - `injection_mode`: The mode of the injection. Default is `4`.
-            - `injection_scale`: The scaling factor for the injection. Default is
-                `1.0`.
+            - `injection_scale`: The scaling factor for the injection. Default
+                is `1.0`.
             - `order`: The order of the Exponential Time Differencing Runge
                 Kutta method. Must be one of {0, 1, 2, 3, 4}. The option `0`
                 only solves the linear part of the equation. Use higher values