feat: detect non-polynomial dependence on the variables

Project.toml

@@ -1,7 +1,7 @@
 name = "QuestBase"
 uuid = "7e80f742-43d6-403d-a9ea-981410111d43"
 authors = ["Orjan Ameye <orjan.ameye@hotmail.com>", "Jan Kosata <kosataj@phys.ethz.ch>", "Javier del Pino <jdelpino@phys.ethz.ch>"]
-version = "0.5.1"
+version = "0.5.2
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

docs/src/API.md

@@ -50,6 +50,8 @@ QuestBase contains a number of symbolic utilities to help with the symbolic mani
 
 ```@docs
 QuestBase.d
+QuestBase.nonpolynomial_terms
+QuestBase.is_polynomial
 ```
 
 ### Exponentials

src/DifferentialEquation.jl

@@ -109,6 +109,24 @@ QuestBase.is_harmonic(diff_eom::DifferentialEquation, t::Num)::Bool =
 """
 $(TYPEDSIGNATURES)
 
+Return the terms of `diff_eom` which are not polynomial in its variables, see
+[`nonpolynomial_terms`](@ref). An empty result means every equation is a polynomial in the
+variables and their derivatives, which averaging requires.
+"""
+function nonpolynomial_terms(diff_eom::DifferentialEquation)
+    vars = get_variables(diff_eom)
+    return unique(
+        reduce(
+            vcat,
+            [nonpolynomial_terms(eq, vars) for eq in values(diff_eom.equations)];
+            init=BasicSymbolic[],
+        ),
+    )
+end
+
+"""
+$(TYPEDSIGNATURES)
+
 Return the independent dependent variables of `diff_eom`.
 """
 function get_independent_variables(diff_eom::DifferentialEquation)

src/Symbolics/Symbolics_utils.jl

@@ -163,6 +163,74 @@ is_harmonic(x, t) = is_harmonic(Num(x), Num(t))
 "Return true if `f` is a function of `var`."
 is_function(f, var) = unwrap(var) in get_variables(f)
 
+"""
+$(TYPEDSIGNATURES)
+
+Return the subexpressions of `x` in which a variable of `vars` (or a derivative of one)
+appears inside an operation other than an addition, a multiplication or a power with a
+non-negative integer exponent. An empty result means `x` is a polynomial in `vars`.
+
+Averaging replaces each variable by a truncated Fourier series and projects the result onto
+the harmonics of the ansatz. Only polynomials map a finite Fourier series onto a finite
+Fourier series, so terms such as `sin(x)`, `exp(x)` or `1/x` cannot be averaged and are
+returned here.
+"""
+function nonpolynomial_terms(x, vars)
+    unwrapped = vars isa AbstractVector ? unwrap.(vars) : [unwrap(vars)]
+    return unique(_collect_nonpolynomial!(BasicSymbolic[], unwrap(x), unwrapped))
+end
+function nonpolynomial_terms(eq::Equation, vars)
+    return unique(
+        cat(nonpolynomial_terms(eq.lhs, vars), nonpolynomial_terms(eq.rhs, vars); dims=1)
+    )
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Return true if `x` is a polynomial in `vars` and their derivatives, i.e. if it has no
+[`nonpolynomial_terms`](@ref).
+"""
+is_polynomial(x, vars) = isempty(nonpolynomial_terms(x, vars))
+
+"Return true if `x` is one of `vars` or contains one of them as a subexpression."
+_contains_variable(x, vars) = any(v -> isequal(x, v), vars)
+function _contains_variable(x::BasicSymbolic, vars)
+    any(v -> isequal(x, v), vars) && return true
+    SymbolicUtils.iscall(x) || return false
+    return any(arg -> _contains_variable(arg, vars), arguments(x))
+end
+
+"Walk `x` and push each subexpression which is not polynomial in `vars` onto `terms`."
+function _collect_nonpolynomial!(terms, x, vars)
+    # a subtree free of `vars` is a coefficient: any operation on it is allowed
+    _contains_variable(x, vars) || return terms
+    any(v -> isequal(x, v), vars) && return terms # a bare variable
+
+    if isadd(x) || ismul(x)
+        for arg in arguments(x)
+            _collect_nonpolynomial!(terms, arg, vars)
+        end
+    elseif ispow(x)
+        base, exponent = arguments(x)
+        power = unwrap_const(exponent)
+        if power isa Real && isinteger(power) && power >= 0
+            _collect_nonpolynomial!(terms, base, vars)
+        else # negative, fractional or symbolic exponent
+            push!(terms, x)
+        end
+    elseif isdiv(x)
+        numerator, denominator = arguments(x)
+        _collect_nonpolynomial!(terms, numerator, vars)
+        _contains_variable(denominator, vars) && push!(terms, x)
+    elseif is_derivative(x)
+        _collect_nonpolynomial!(terms, first(arguments(x)), vars)
+    else # a call such as sin, exp or log applied to something containing a variable
+        push!(terms, x)
+    end
+    return terms
+end
+
 """
 Counts the number of derivatives of a symbolic variable.
 """

test/symbolics.jl

@@ -103,6 +103,40 @@ end
     @test !is_harmonic(dEOM, t)
 end
 
+@testset "polynomial in the variables" begin
+    using QuestBase: is_polynomial, nonpolynomial_terms, d, DifferentialEquation
+
+    @variables a, b, t, x(t), y(t), ω
+
+    @test is_polynomial(a * x^3 + b * d(x, t) * x^2 - a, [x])
+    @test is_polynomial(a * cos(ω * t) * x - b * x * y, [x, y])
+    @test is_polynomial(d(x, t, 2) + a * x / b, [x]) # variable-free denominator
+    @test is_polynomial(a + cos(ω * t), [x]) # no variable at all
+    @test is_polynomial(a * x^2.0, [x]) # integer-valued exponent of another type
+    @test all(n -> is_polynomial(a * x^n + b * d(x, t) * x^n, [x]), 0:12)
+
+    @test !is_polynomial(a * sin(x), [x])
+    @test !is_polynomial(a * exp(x + y), [x, y])
+    @test !is_polynomial(a / x, [x])
+    @test !is_polynomial(x^(-2), [x])
+    @test !is_polynomial(sqrt(x), [x])
+    @test !is_polynomial(a^x, [x])
+    @test !is_polynomial(d(sin(x), t), [x]) # hidden inside a derivative
+
+    # the offending terms are reported, the harmless ones are not
+    terms = nonpolynomial_terms(a * x^3 + a * sin(x) ~ b * cos(ω * t), [x])
+    @eqtest Num.(terms) == [sin(x)]
+
+    # a variable of another equation does not make a term non-polynomial
+    @test is_polynomial(sin(y) + x, [x])
+
+    dEOM = DifferentialEquation(
+        [d(x, t, 2) + x + a * x^3, d(y, t, 2) + y + a * sin(y)], [x, y]
+    )
+    @eqtest Num.(nonpolynomial_terms(dEOM)) == [sin(y)]
+    @test isempty(nonpolynomial_terms(DifferentialEquation([x + a * x^3, y], [x, y])))
+end
+
 @testset "fourier" begin
     using QuestBase: fourier_cos_term, fourier_sin_term
     using QuestBase.Symbolics: expand