Implement FourierBandForcing (#32)

* Allow copying coefficients from Fourier field

* Implement FourierBandForcing (untested)

* Update Timestepping module

* Small fixes + add tests

* Improve from_fourier_grid!

* Some fixes + tests

* Update CHANGELOG [skip ci]

* Update test

* Fix nonperiodic tests

* Allow α′ term in FourierBandForcing

* Add option to use coarse-grained vorticity field

* Simpler variant of coarse-grained vorticity case

* Treat case where coarse-grained vorticity is zero

* Test with smooth_vorticity = true

* smoothed_vorticity -> filtered_vorticity

* Update docs

* Add SyntheticFields.remove_zeros!

* Fixes to docs

* Parallelise forcing computations

Author: jipolanco

CHANGELOG.md

@@ -5,6 +5,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Implement `FourierBandForcing` as an alternative to `NormalFluidForcing`.
+  The new forcing is expected to influence only (or mostly) the scales where it
+  is defined, letting the other lengthscales "free" to evolve.
+
 ## [0.26.4] - 2025-03-04
 
 ### Changed

Project.toml

@@ -19,6 +19,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonuniformFFTs = "cd96f58b-6017-4a02-bb9e-f4d81626177f"
+OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -62,6 +63,7 @@ LinearAlgebra = "1.9"
 MakieCore = "0.6, 0.7, 0.8, 0.9"
 NonuniformFFTs = "0.5.1, 0.6, 0.7"
 Observables = "0.5"
+OhMyThreads = "0.7.0"
 Polyester = "0.7"
 PrecompileTools = "1"
 Random = "1.10"

docs/src/modules/Forcing.md

@@ -14,7 +14,12 @@ Forcing
 ```@docs
 NormalFluidForcing
 apply!
-get_velocities!
+```
+
+## Fourier band forcing
+
+```@docs
+FourierBandForcing
 ```
 
 ## Abstract types

docs/src/modules/SyntheticFields.md

@@ -20,6 +20,8 @@ UniformVectorField
 ```@docs
 FourierBandVectorField
 init_coefficients!
+from_fourier_grid!
+remove_zeros!
 ```
 
 ## Saving and loading data

src/Forcing/Forcing.jl

@@ -6,10 +6,14 @@ Defines methods for injecting energy onto a system of vortices.
 module Forcing
 
 using ..Filaments: AbstractFilament, UnitTangent
-using ..SyntheticFields: SyntheticFields  # for docs only
+using ..BiotSavart: BiotSavart, BiotSavartCache
+using ..SyntheticFields: SyntheticFields, FourierBandVectorField
 using LinearAlgebra: ×
+using OhMyThreads: Scheduler, SerialScheduler, tforeach
 
-export AbstractForcing, NormalFluidForcing
+using Adapt: adapt
+
+export AbstractForcing, NormalFluidForcing, FourierBandForcing
 
 """
     AbstractForcing
@@ -18,157 +22,28 @@ Abstract type representing a forcing method.
 """
 abstract type AbstractForcing end
 
+init_cache(forcing::Nothing, args...) = nothing  # called when forcing is disabled
+
 """
-    Forcing.apply!(forcing::AbstractForcing, vs::AbstractVector{<:Vec3}, f::AbstractFilament)
+    Forcing.apply!(forcing::AbstractForcing, vs::AbstractVector{<:Vec3}, f::AbstractFilament; [scheduler])
 
 Apply forcing to a single filament `f` with self-induced velocities `vs`.
 
 At output, the `vs` vector is overwritten with the actual vortex line velocities.
 
+The optional `scheduler` keyword can be used to parallelise computations using one of the
+[schedulers defined in OhMyThreads.jl](https://juliafolds2.github.io/OhMyThreads.jl/stable/refs/api/#Schedulers).
+
 ---
 
-    Forcing.apply!(forcing::NormalFluidForcing, vs, vn, tangents)
+    Forcing.apply!(forcing::NormalFluidForcing, vs, vn, tangents; [scheduler])
 
 This variant can be used in the case of a [`NormalFluidForcing`](@ref) if one already has
 precomputed values of the normal fluid velocity and local unit tangents at filament points.
-
-The normal fluid velocities may be computed using [`Forcing.get_velocities!`](@ref).
 """
 function apply! end
 
-@doc raw"""
-    NormalFluidForcing <: AbstractForcing
-    NormalFluidForcing(vn::Function; α, α′ = 0)
-
-Forcing due to mutual friction with a normal fluid.
-
-The normal fluid is represented by a function `vn` which should take a position
-`x⃗::SVector{N, T}` and return a velocity `v⃗::SVector{N, T}` (`N` is the number of
-dimensions, usually `N = 3`).
-
-In particular, the function could be a synthetic velocity field from the
-[`SyntheticFields`](@ref) module (see below for examples).
-
-This type of forcing defines an external velocity ``\bm{v}_{\text{f}}`` affecting vortex
-motion, so that the actual vortex velocity ``\bm{v}_{\text{L}}`` is
-
-```math
-\frac{\mathrm{d}\bm{s}}{\mathrm{d}t} = \bm{v}_{\text{L}} = \bm{v}_{\text{s}} + \bm{v}_{\text{f}}.
-```
-
-Here ``\bm{v}_{\text{s}}`` is the self-induced velocity obtained by applying Biot–Savart's law.
-
-The forcing velocity is of the form:
-
-```math
-\bm{v}_{\text{f}} =
-α \bm{s}' × (\bm{v}_{\text{n}} - \bm{v}_{\text{s}})
-- α' \bm{s}' × \left[ \bm{s}' × (\bm{v}_{\text{n}} - \bm{v}_{\text{s}}) \right]
-```
-
-where ``\bm{s}'`` is the local unit tangent vector, and ``α`` and ``α'`` are non-dimensional
-coefficients representing the intensity of Magnus and drag forces.
-
-# Example
-
-Define a mutual friction forcing based on a large-scale normal fluid velocity field
-(see [`SyntheticFields.FourierBandVectorField`](@ref)):
-
-```jldoctest
-julia> using VortexPasta.Forcing: NormalFluidForcing
-
-julia> using VortexPasta.SyntheticFields: SyntheticFields, FourierBandVectorField
-
-julia> using Random: Xoshiro
-
-julia> rng = Xoshiro(42);  # initialise random number generator (optional, but recommended)
-
-julia> Ls = (2π, 2π, 2π);  # domain dimensions
-
-julia> vn_rms = 1.0;  # typical magnitude (rms value) of normal fluid velocity components
-
-julia> vn = FourierBandVectorField(undef, Ls; kmin = 0.1, kmax = 1.5)  # create field with non-zero Fourier wavevectors kmin ≤ |k⃗| ≤ kmax
-FourierBandVectorField{Float64, 3} with 9 independent Fourier coefficients in |k⃗| ∈ [1.0, 1.4142]
-
-julia> SyntheticFields.init_coefficients!(rng, vn, vn_rms);  # randomly set non-zero Fourier coefficients of the velocity field
-
-julia> forcing = NormalFluidForcing(vn; α = 0.8, α′ = 0)
-NormalFluidForcing{Float64} with:
- ├─ Magnus force coefficient: α = 0.8
- ├─ Drag force coefficient: α′ = 0.0
- └─ Normal velocity field: FourierBandVectorField{Float64, 3} with 9 independent Fourier coefficients in |k⃗| ∈ [1.0, 1.4142]
-```
-
-"""
-struct NormalFluidForcing{
-        T,
-        VelocityField <: Function,  # should return an SVector{N, T}
-    } <: AbstractForcing
-    vn :: VelocityField
-    α  :: T
-    α′ :: T
-end
-
-function NormalFluidForcing(vn::F; α::T, α′::Real = 0) where {T <: AbstractFloat, F <: Function}
-    NormalFluidForcing(vn, T(α), T(α′))
-end
-
-function Base.show(io::IO, f::NormalFluidForcing{T}) where {T}
-    (; vn, α, α′,) = f
-    indent = get(io, :indent, 0)
-    nspaces = max(indent, 1)
-    spaces = " "^nspaces
-    print(io, "NormalFluidForcing{$T} with:")
-    print(io, "\n$(spaces)├─ Magnus force coefficient: α = ", α)
-    print(io, "\n$(spaces)├─ Drag force coefficient: α′ = ", α′)
-    print(io, "\n$(spaces)└─ Normal velocity field: ", vn)
-end
-
-"""
-    Forcing.get_velocities!(forcing::NormalFluidForcing, vn::AbstractVector{<:Vec3}, f::AbstractFilament)
-
-Evaluate normal fluid velocity at filament nodes.
-
-Results are written to `vn`, which should have the same length as the filament `f`.
-"""
-function get_velocities!(forcing::NormalFluidForcing, vn::AbstractVector, f::AbstractFilament)
-    for i ∈ eachindex(vn, f)
-        vn[i] = forcing.vn(f[i])
-    end
-    vn
-end
-
-function apply!(forcing::NormalFluidForcing, vs::AbstractVector, f::AbstractFilament)
-    eachindex(vs) == eachindex(f) || throw(DimensionMismatch("lengths of filament and velocity vectors don't match"))
-    @inline get_at_node(i) = (v⃗ₙ = forcing.vn(f[i]), s⃗′ = f[i, UnitTangent()])
-    _apply!(get_at_node, forcing, vs)
-end
-
-function apply!(forcing::NormalFluidForcing, vs::AbstractVector, vn::AbstractVector, tangents::AbstractVector)
-    eachindex(vs) == eachindex(vn) == eachindex(tangents) || throw(DimensionMismatch("lengths of vectors don't match"))
-    @inline get_at_node(i) = @inbounds (v⃗ₙ = vn[i], s⃗′ = tangents[i],)
-    _apply!(get_at_node, forcing, vs)
-end
-
-# The first argument is a `get_at_node(i)` function which returns a NamedTuple with fields
-# (v⃗ₛ, v⃗ₙ, s⃗′) with the superfluid velocity, normal fluid velocity and the local unit
-# tangent at the node `i` of a filament. This is useful if those quantities have been
-# precomputed.
-function _apply!(get_at_node::F, forcing::NormalFluidForcing, vs::AbstractVector) where {F <: Function}
-    (; α, α′,) = forcing
-    V = eltype(vs)  # usually Vec3{T} = SVector{3, T}
-    for i ∈ eachindex(vs)
-        (; v⃗ₙ, s⃗′,) = @inline get_at_node(i)  # superfluid velocity, normal fluid velocity and unit tangent
-        v⃗ₛ = vs[i]
-        vₙₛ = V(v⃗ₙ) - V(v⃗ₛ)   # slip velocity
-        v_perp = s⃗′ × vₙₛ
-        vf = α * v_perp    # velocity due to Magnus force
-        if !iszero(α′)
-            vf = vf - α′ * s⃗′ × v_perp  # velocity due to drag force (it's quite common to set α′ = 0)
-        end
-        vs[i] = v⃗ₛ + vf
-    end
-    vs
-end
+include("normal_fluid.jl")
+include("fourier_band.jl")
 
 end

src/Forcing/fourier_band.jl

@@ -0,0 +1,158 @@
+@doc raw"""
+    FourierBandForcing <: AbstractForcing
+    FourierBandForcing(vn::FourierBandVectorField; α, α′ = 0, filtered_vorticity = false)
+
+Forcing due to mutual friction of a normal fluid with a Fourier-filtered superfluid velocity.
+
+This forcing is similar to [`NormalFluidForcing`](@ref), but tries to only affect scales
+within a given band `[kmin, kmax]` in Fourier space. This is achieved by a normal fluid velocity
+field represented by a [`FourierBandVectorField`](@ref), and by a modified Schwarz's
+equation in which only a coarse-grained superfluid flow is taken into account in the
+estimation of the mutual friction term.
+
+Concretely, the vortex line velocity according to this forcing type is:
+
+```math
+\frac{\mathrm{d}\bm{s}}{\mathrm{d}t} = \bm{v}_{\text{L}} = \bm{v}_{\text{s}} + \bm{v}_{\text{f}}
+```
+
+The forcing velocity is of the form:
+
+```math
+\bm{v}_{\text{f}} = α \bm{s}' × \bm{v}_{\text{ns}}^{>} - α' \bm{s}' × \left( \bm{s}' × \bm{v}_{\text{ns}}^{>} \right)
+```
+
+where ``\bm{v}_{\text{ns}}^{>} = \bm{v}_{\text{n}} - \bm{v}_{\text{s}}^{>}`` is a filtered slip velocity.
+In practice, the filtered velocity is active within the same wavenumber range `[kmin, kmax]` where `vn` is
+defined. See [`NormalFluidForcing`](@ref) for other definitions.
+
+## Using a filtered vorticity field (experimental)
+
+To further ensure that this forcing only affects the chosen range of scales, one can pass
+`filtered_vorticity = true`, which will replace the local unit tangent ``\bm{s}'`` with a
+normalised coarse-grained vorticity. This corresponds to setting ``\bm{s}' = \bm{ω}^{>} / |\bm{ω}^{>}|``
+where ``\bm{ω}^{>}`` is the Fourier-filtered vorticity field.
+"""
+struct FourierBandForcing{
+        T <: AbstractFloat,
+        N,  # number of dimensions (usually 3)
+        VelocityField <: FourierBandVectorField{T, N},
+    } <: AbstractForcing
+    vn :: VelocityField
+    α  :: T
+    α′ :: T
+    filtered_vorticity :: Bool
+end
+
+with_filtered_vorticity(f::FourierBandForcing) = f.filtered_vorticity
+
+function FourierBandForcing(
+        vn::FourierBandVectorField{T, 3}; α::Real, α′::Real = 0,
+        filtered_vorticity::Bool = false,
+    ) where {T <: AbstractFloat}
+    FourierBandForcing(vn, T(α), T(α′), filtered_vorticity)
+end
+
+function Base.show(io::IO, f::FourierBandForcing{T}) where {T}
+    (; vn, α, α′,) = f
+    indent = get(io, :indent, 0)
+    nspaces = max(indent, 1)
+    spaces = " "^nspaces
+    print(io, "FourierBandForcing{$T} with:")
+    print(io, "\n$(spaces)├─ Normal velocity field: ", vn)
+    print(io, "\n$(spaces)├─ Filtered superfluid vorticity: ", with_filtered_vorticity(f))
+    print(io, "\n$(spaces)└─ Friction coefficients: α = ", α, " and α′ = ", α′)
+end
+
+# Here vs_d is the superfluid velocity in Fourier space, optionally on a device (GPU).
+function init_cache(f::FourierBandForcing{T, N}, cache_bs::BiotSavartCache) where {T, N}
+    (; vn,) = f
+    vs_d = BiotSavart.get_longrange_field_fourier(cache_bs).field
+    A = eltype(vs_d).name.wrapper  # e.g. Array, CuArray, ... (XXX: relies on Julia internals!)
+    vn_d = adapt(A, vn)::FourierBandVectorField        # vn on the device
+    vtmp_h = similar(vn)::FourierBandVectorField       # temporary buffer (on host)
+    vtmp_d = adapt(A, vtmp_h)::FourierBandVectorField  # temporary buffer (on device)
+    ω_h = similar(vtmp_h)  # coarse-grained superfluid vorticity
+    if !with_filtered_vorticity(f)
+        empty!(ω_h)  # we don't use this field
+    end
+    ω_d = adapt(A, ω_h)
+    (; vn_d, vtmp_d, vtmp_h, ω_h, ω_d)
+end
+
+function update_cache!(cache, f::FourierBandForcing{T, N}, cache_bs::BiotSavartCache) where {T, N}
+    (; vtmp_d, vn_d, vtmp_h, ω_h, ω_d) = cache
+
+    vs_d, ks_grid = let data = BiotSavart.get_longrange_field_fourier(cache_bs)
+        local (; state, field, wavenumbers,) = data
+        @assert state.quantity == :velocity
+        @assert state.smoothed == true
+        field, wavenumbers
+    end
+    α_ewald = cache_bs.params.α
+
+    # (1) Copy normal fluid velocity onto buffer (device -> device copy)
+    @assert vtmp_d.cs !== vn_d.cs  # coefficients are not aliased
+    copyto!(vtmp_d, vn_d)
+
+    # (2) Retrieve values from coarse-grained velocity field (Gaussian filtered with α_ewald parameter).
+    # We "unfilter" the values, similarly to the way we compute energy spectra from the long-range velocity.
+    inv_four_α² = 1 / (4 * α_ewald * α_ewald)
+    @inline function op(vn, vs_filtered, k⃗)
+        k² = sum(abs2, k⃗)
+        φ = @fastmath exp(k² * inv_four_α²)
+        vs = φ * vs_filtered
+        vn - vs
+    end
+    SyntheticFields.from_fourier_grid!(op, vtmp_d, vs_d, ks_grid)  # vtmp_d now contains vn_d(k⃗) - vs_d(k⃗) in [kmin, kmax]
+
+    # Optionally compute coarse-grained vorticity.
+    @inline function op_vorticity(_, vs_filtered, k⃗)
+        k² = sum(abs2, k⃗)
+        φ = @fastmath exp(k² * inv_four_α²)
+        vs = φ * vs_filtered
+        im * (k⃗ × vs)
+    end
+    if with_filtered_vorticity(f)
+        @assert length(ω_h) == length(ω_d) == length(vtmp_d)
+        SyntheticFields.from_fourier_grid!(op_vorticity, ω_d, vs_d, ks_grid)
+    end
+
+    # (3) Copy results to CPU if needed (avoided if the "device" is the CPU).
+    if vtmp_h !== vtmp_d
+        copyto!(vtmp_h, vtmp_d)
+    end
+    if ω_h !== ω_d && with_filtered_vorticity(f)
+        copyto!(ω_h, ω_d)
+    end
+
+    nothing
+end
+
+function apply!(forcing::FourierBandForcing, cache, vs::AbstractVector, f::AbstractFilament; scheduler = SerialScheduler())
+    (; α, α′,) = forcing
+    (; vtmp_h, ω_h,) = cache  # contains vₙ - vₛ at large scale
+    V = eltype(vs)  # usually Vec3{T} = SVector{3, T}
+    tforeach(eachindex(vs); scheduler) do i
+        @inline
+        s⃗ = f[i]
+        if with_filtered_vorticity(forcing)
+            ω⃗ = V(ω_h(s⃗))  # coarse-grained vorticity at s⃗
+            ω_norm = sqrt(sum(abs2, ω⃗))
+            # Note: in the case that the coarse-grained vorticity is exactly zero (very very
+            # unlikely), we simply set s⃗′ to zero which disables the forcing.
+            ω_invnorm = ifelse(iszero(ω_norm), ω_norm, 1/ω_norm)  # = 1/|ω⃗| (or zero, if |ω⃗| = 0)
+            s⃗′ = (ω⃗ * ω_invnorm)::V  # replace s⃗′ with direction of coarse-grained vorticity
+        else
+            s⃗′ = f[i, UnitTangent()]::V
+        end
+        v⃗ₙₛ = V(vtmp_h(s⃗))
+        v_perp = s⃗′ × v⃗ₙₛ
+        vf = α * v_perp
+        if !iszero(α′)
+            vf = vf - α′ * s⃗′ × v_perp
+        end
+        vs[i] = vs[i] + vf
+    end
+    vs
+end