feat: add additional keywords to ia_solve

Author: AayushSabharwal

src/solver/ia_main.jl

@@ -2,7 +2,7 @@ using Symbolics
 
 const SAFE_ALTERNATIVES = Dict(log => slog, sqrt => ssqrt, cbrt => scbrt)
 
-function isolate(lhs, var; warns=true, conditions=[])
+function isolate(lhs, var; warns=true, conditions=[], complex_roots = true, periodic_roots = true)
     rhs = Vector{Any}([0])
     original_lhs = deepcopy(lhs)
     lhs = unwrap(lhs)
@@ -74,12 +74,21 @@ function isolate(lhs, var; warns=true, conditions=[])
                 power = args[2]
                 new_roots = []
 
-                for i in eachindex(rhs)
-                    for k in 0:(args[2] - 1)
-                        r = wrap(term(^, rhs[i], (1 // power)))
-                        c = wrap(term(*, 2 * (k), pi)) * im / power
-                        root = r * Base.MathConstants.e^c
-                        push!(new_roots, root)
+                if complex_roots
+                    for i in eachindex(rhs)
+                        for k in 0:(args[2] - 1)
+                            r = wrap(term(^, rhs[i], (1 // power)))
+                            c = wrap(term(*, 2 * (k), pi)) * im / power
+                            root = r * Base.MathConstants.e^c
+                            push!(new_roots, root)
+                        end
+                    end
+                else
+                    for i in eachindex(rhs)
+                        push!(new_roots, wrap(term(^, rhs[i], (1 // power))))
+                        if iseven(power)
+                            push!(new_roots, wrap(term(-, new_roots[end])))
+                        end
                     end
                 end
                 rhs = []
@@ -97,7 +106,7 @@ function isolate(lhs, var; warns=true, conditions=[])
             ia_conditions!(oper, lhs, rhs, conditions)
             invop = left_inverse(oper)
             invop = get(SAFE_ALTERNATIVES, invop, invop)
-            if is_periodic(oper)
+            if is_periodic(oper) && periodic_roots
                 # make this global somehow so the user doesnt need to declare it on his own
                 new_var = gensym()
                 new_var = (@variables $new_var)[1]
@@ -118,7 +127,7 @@ function isolate(lhs, var; warns=true, conditions=[])
     return rhs, conditions
 end
 
-function attract(lhs, var; warns = true)
+function attract(lhs, var; warns = true, complex_roots = true, periodic_roots = true)
     if n_func_occ(simplify(lhs), var) <= n_func_occ(lhs, var)
         lhs = simplify(lhs)
     end
@@ -133,7 +142,9 @@ function attract(lhs, var; warns = true)
     end
     lhs = attract_trig(lhs, var)
 
-    n_func_occ(lhs, var) == 1 && return isolate(lhs, var, warns = warns, conditions=conditions)
+    if n_func_occ(lhs, var) == 1
+        return isolate(lhs, var; warns, conditions, complex_roots, periodic_roots)
+    end
 
     lhs, sub = turn_to_poly(lhs, var)
 
@@ -151,12 +162,12 @@ function attract(lhs, var; warns = true)
     new_var = collect(keys(sub))[1]
     new_var_val = collect(values(sub))[1]
 
-    roots, new_conds = isolate(lhs, new_var, warns = warns)
+    roots, new_conds = isolate(lhs, new_var; warns = warns, complex_roots, periodic_roots)
     append!(conditions, new_conds)
     new_roots = []
 
     for root in roots
-        new_sol, new_conds = isolate(new_var_val - root, var, warns = warns)
+        new_sol, new_conds = isolate(new_var_val - root, var; warns = warns, complex_roots, periodic_roots)
         append!(conditions, new_conds)
         push!(new_roots, new_sol)
     end
@@ -166,7 +177,7 @@ function attract(lhs, var; warns = true)
 end
 
 """
-    ia_solve(lhs, var)
+    ia_solve(lhs, var; kwargs...)
 This function attempts to solve transcendental functions by first checking
 the "smart" number of occurrences in the input LHS. By smart here we mean
 that polynomials are counted as 1 occurrence. for example `x^2 + 2x` is 1
@@ -195,6 +206,13 @@ we throw an error to tell the user that this is currently unsolvable by our cove
 - lhs: a Num/SymbolicUtils.BasicSymbolic
 - var: variable to solve for.
 
+# Keyword arguments
+- `warns = true`: Whether to emit warnings for unsolvable expressions.
+- `complex_roots = true`: Whether to consider complex roots of `x ^ n ~ y`, where `n` is an integer.
+- `periodic_roots = true`: If `true`, isolate `f(x) ~ y` as `x ~ finv(y) + n * period` where
+   `is_periodic(f) == true`, `finv = left_inverse(f)` and `period = fundamental_period(f)`. `n`
+   is a new anonymous symbolic variable.
+
 # Examples
 ```jldoctest
 julia> solve(a*x^b + c, x)
@@ -238,17 +256,17 @@ See also: [`left_inverse`](@ref), [`inverse`](@ref), [`is_periodic`](@ref),
 # References
 [^1]: [R. W. Hamming, Coding and Information Theory, ScienceDirect, 1980](https://www.sciencedirect.com/science/article/pii/S0747717189800070).
 """
-function ia_solve(lhs, var; warns = true)
+function ia_solve(lhs, var; warns = true, complex_roots = true, periodic_roots = true)
     nx = n_func_occ(lhs, var)
     sols = []
     conditions = []
     if nx == 0
         warns && @warn("Var not present in given expression")
         return []
     elseif nx == 1
-        sols, conditions = isolate(lhs, var, warns = warns)
+        sols, conditions = isolate(lhs, var; warns = warns, complex_roots, periodic_roots)
     elseif nx > 1
-        sols, conditions = attract(lhs, var, warns = warns)
+        sols, conditions = attract(lhs, var; warns = warns, complex_roots, periodic_roots)
     end
 
     isequal(sols, nothing) && return nothing