@@ -35,11 +35,12 @@ mutable struct BookKeeping
3535 v_array4 :: Vector{Symbol}
3636 v_preamb :: Vector{Union{Symbol,Expr}}
3737 retvar :: Symbol
38+ numaux :: Int
3839
3940 function BookKeeping ()
4041 return new (Dict {Symbol, Expr} (), Dict {Union{Symbol,Expr}, Number} (),
4142 Dict {Symbol, Expr} (), Symbol[], Symbol[], Symbol[], Symbol[], Symbol[], Symbol[],
42- Union{Symbol,Expr}[], :nothing )
43+ Union{Symbol,Expr}[], :nothing , 0 )
4344 end
4445end
4546
@@ -123,7 +124,7 @@ const _HEAD_ALLOC_TAYLOR1_VECTOR = sanitize(:(
123124# Constants for the initial declaration and initialization of arrays
124125const _DECL_ARRAY = sanitize ( Expr (:block ,
125126 :(__var1 = Array {Taylor1{_S}} (undef, __var2)),
126- :( for i in CartesianIndices (__var1) __var1[i] = Taylor1 ( zero (constant_term (__x[1 ])), order ) end ))
127+ :( for i in eachindex (__var1) __var1[i] = Taylor1 ( zero (constant_term (__x[1 ])), order ) end ))
127128);
128129
129130
@@ -147,25 +148,25 @@ function _make_parsed_jetcoeffs(ex::Expr)
147148 new_jetcoeffs, new_allocjetcoeffs = _newhead (fn, fnargs)
148149
149150 # Transform the graph representation of the body of the functions:
150- # defspreamble: inicializations used for the zeroth order (preamble)
151+ # defspreamble: initializations used for the zeroth order (preamble)
151152 # defsprealloc: definitions (declarations) of auxiliary Taylor1's
152153 # fnbody: transformed function body, using mutating functions from TaylorSeries;
153154 # used later within the recursion loop
154155 # bkkeep: book-keeping structure having info of the variables
155156 defspreamble, defsprealloc, fnbody, bkkeep = _preamble_body (fnbody, fnargs)
156157
157158 # Create body of recursion loop; temporary assignements may be needed.
158- # rec_preamb: recursion loop for the preamble (first order correction)
159- # rec_fnbody: recursion loop for the body-function (recursion loop for higher orders)
160- rec_preamb, rec_fnbody = _recursionloop (fnargs, bkkeep)
159+ # rec_fnbody: recursion loop for the body-function (recursion loop for all orders)
160+ rec_fnbody = _recursionloop (fnargs, bkkeep)
161161
162162 # Expr for the for-loop block for the recursion (of the `x` variable)
163- forloopblock = Expr (:for , :(ord = 1 : order- 1 ), Expr (:block , :(ordnext = ord + 1 )) )
164- # Add rec_fnbody to forloopblock
163+ forloopblock = Expr (:for , :(ord = 0 : order- 1 ), Expr (:block , :(ordnext = ord + 1 )) )
164+
165+ # Add rec_fnbody to `forloopblock`
165166 push! (forloopblock. args[2 ]. args, fnbody. args[1 ]. args... , rec_fnbody)
166167
167168 # Add preamble and recursion body to `new_jetcoeffs`
168- push! (new_jetcoeffs. args[2 ]. args, defspreamble... , rec_preamb )
169+ push! (new_jetcoeffs. args[2 ]. args, defspreamble... )
169170
170171 # Push preamble and forloopblock to `new_jetcoeffs` and return line
171172 push! (new_jetcoeffs. args[2 ]. args, forloopblock, Meta. parse (" return nothing" ))
@@ -337,11 +338,11 @@ of the original diferential equations function.
337338
338339"""
339340function _preamble_body (fnbody, fnargs)
340- # Inicialize BookKeeping struct
341+ # Initialize BookKeeping struct
341342 bkkeep = BookKeeping ()
342343
343344 # Rename vars to have the body in non-indexed form; bkkeep has different entries
344- # for bookkeeping variables/symbolds , including indexed ones
345+ # for bookkeeping variables/symbols , including indexed ones
345346 fnbody, bkkeep. d_indx = _rename_indexedvars (fnbody)
346347
347348 # Create `newfnbody` which corresponds to `fnbody`, cleaned (without irrelevant comments)
@@ -362,15 +363,22 @@ function _preamble_body(fnbody, fnargs)
362363 preamble = subs (preamble, bkkeep. d_assign)
363364 prealloc = subs (prealloc, bkkeep. d_assign)
364365
365- # Include the assignement of indexed auxiliary variables
366+ # Include the assignment of indexed auxiliary variables
366367 defsprealloc = _defs_allocs! (prealloc, fnargs, bkkeep, false )
367368 preamble = subs (preamble, bkkeep. d_indx)
368- defspreamble = Expr[preamble. args... ]
369+
370+ # Retain only `local` declarations from `preamble` in `new_preamble`
371+ new_preamble = Expr (:block ,)
372+ for (i, arg) in enumerate (preamble. args)
373+ (arg. head == :local ) && push! (new_preamble. args, arg)
374+ end
375+ defspreamble = Expr[new_preamble. args... ]
376+
369377 # Bring back substitutions
370378 newfnbody = subs (newfnbody, bkkeep. d_indx)
371379
372380 # Define retvar; for scalar eqs is the last entry included in v_newvars
373- bkkeep. retvar = length (fnargs) == 3 ? subs (bkkeep. v_newvars[end ], bkkeep. d_indx) : fnargs[1 ]
381+ bkkeep. retvar = length (fnargs) == 3 ? subs (bkkeep. v_newvars[end - bkkeep . numaux ], bkkeep. d_indx) : fnargs[1 ]
374382
375383 return defspreamble, defsprealloc, newfnbody, bkkeep
376384end
@@ -834,7 +842,7 @@ function _replacecalls!(bkkeep::BookKeeping, fnold::Expr, newvar::Symbol)
834842 end
835843
836844 elseif ll == 3
837- # Binary call; no auxiliary expressions needed
845+ # Binary call
838846 newarg2 = fnold. args[3 ]
839847
840848 # Replacements
@@ -859,6 +867,28 @@ function _replacecalls!(bkkeep::BookKeeping, fnold::Expr, newvar::Symbol)
859867 # Dict(:_res => newvar, :_arg1 => :(constant_term($(newarg1))),
860868 # :_arg2 => :(constant_term($(newarg2))), :_k => :ord))
861869
870+ # Auxiliary expression
871+ if aux_fnexpr. head != :nothing
872+ newaux = genname ()
873+
874+ aux_alloc = :( _res = Taylor1 ($ (aux_fnexpr. args[2 ]), order) )
875+ aux_alloc = subs (aux_alloc,
876+ Dict (:_res => newaux, :_arg1 => :(constant_term ($ (newarg1))), :_aux => newaux))
877+
878+ aux_fnexpr = Expr (:block ,
879+ :(TaylorSeries. zero! (_res)),
880+ :(_res. coeffs[1 ] = $ (aux_fnexpr. args[2 ])) )
881+ aux_fnexpr = subs (aux_fnexpr,
882+ Dict (:_res => newaux, :_arg1 => :(constant_term ($ (newarg1))), :_aux => newaux))
883+
884+ fnexpr = subs (fnexpr, Dict (:_aux => newaux))
885+ if newvar ∈ bkkeep. v_arraydecl
886+ push! (bkkeep. v_arraydecl, newaux)
887+ else
888+ bkkeep. numaux += 1
889+ push! (bkkeep. v_newvars, newaux)
890+ end
891+ end
862892 else
863893 # Recognized call, but not a unary or binary call; copy expression
864894 fnexpr = :($ newvar = $ fnold)
@@ -1037,27 +1067,21 @@ function _recursionloop(fnargs, bkkeep::BookKeeping)
10371067 ll = length (fnargs)
10381068
10391069 if ll == 3
1040- rec_preamb = sanitize ( :( $ (fnargs[1 ]). coeffs[2 ] = $ (bkkeep. retvar). coeffs[1 ] ) )
1041- rec_fnbody = sanitize ( :( $ (fnargs[1 ]). coeffs[ordnext+ 1 ] = $ (bkkeep. retvar). coeffs[ordnext]/ ordnext ) )
1070+ rec_fnbody = sanitize ( :( TaylorIntegration. solcoeff! ($ (fnargs[1 ]), $ (bkkeep. retvar), ordnext) ) )
10421071
10431072 elseif ll == 4
10441073 bkkeep. retvar = fnargs[1 ]
1045- rec_preamb = sanitize (:(
1046- for __idx in eachindex ($ (fnargs[2 ]))
1047- $ (fnargs[2 ])[__idx]. coeffs[2 ] = $ (bkkeep. retvar)[__idx]. coeffs[1 ]
1048- end ))
10491074 rec_fnbody = sanitize (:(
10501075 for __idx in eachindex ($ (fnargs[2 ]))
1051- $ (fnargs[2 ])[__idx]. coeffs[ordnext+ 1 ] =
1052- $ (bkkeep. retvar)[__idx]. coeffs[ordnext]/ ordnext
1076+ TaylorIntegration. solcoeff! ($ (fnargs[2 ])[__idx], $ (bkkeep. retvar)[__idx], ordnext)
10531077 end ))
10541078
10551079 else
10561080 throw (ArgumentError (
10571081 " Wrong number of arguments in the definition of the function $fn " ))
10581082 end
10591083
1060- return rec_preamb, rec_fnbody
1084+ return rec_fnbody
10611085end
10621086
10631087
0 commit comments