Add functionality; backward integration (#4)

Author: JinraeKim

Project.toml

@@ -1,7 +1,7 @@
 name = "FymEnvs"
 uuid = "d6fd7ba0-2ca9-4676-ba23-5bed9e863cfb"
 authors = ["JinraeKim <kjl950403@gmail.com> and contributors"]
-version = "0.2.4"
+version = "0.3.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"

README.md

@@ -47,7 +47,7 @@ systems = Dict("sys" => BaseSystem(initial_state=x0, name="3d_sys"))
 log_dir = "data/test"
 file_name = "fym.h5"
 logger = Logger(log_dir=log_dir, file_name=file_name)
-env = BaseEnv(max_t=100.00, logger=logger, name="test_env")
+env = BaseEnv(max_t=100.00, logger=logger, name="test_env",)
 systems!(env, systems)  # set systems; required
 dyn!(env, set_dyn)  # set dynamics; required
 step!(env, step)  # set step; required
@@ -66,27 +66,27 @@ i = 0
 end
 close!(env)
 data = load(env.logger.path)
-@show env
-@show size(data["state"]["sys"])
+show(env)
+show(size(data["state"]["sys"]))
 ```
 
 Result:
 
 ```julia
 # time and progressbar
-99%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌ |  ETA: 0:00:00
-0.744015 seconds (3.85 M allocations: 225.629 MiB, 3.30% gc time)
+ 99%|████████████████████████████████████████████████████████████████████████████████████▏|  ETA: 0:00:00  3.513103 seconds (11.45 M allocations: 621.034 MiB, 3.28% gc time)
 # representation, i.e., show (nested env supported)
 name: test_env
+max_t: 100.0
+dt: 0.01
 +---name: 3d_sys
 |   state: [3.7200760072278747e-44, 7.440152014455749e-44, 1.1160228021683672e-43]
 |   dot: [-3.7200756925403154e-44, -7.440151385080631e-44, -1.1160227077620995e-43]
 |   initial_state: [1.0, 2.0, 3.0]
 |   state_size: (3,)
 |   flat_index: 1:3
-env =
 # saved data
-size((data["state"])["sys"]) = (10000, 3)
+(10000, 3)
 ```
 
 For more examples, see directory `test`.

src/FymCore.jl

@@ -31,6 +31,13 @@ export time_over, time
 
 
 ############ Clock ############
+"""
+    Clock
+
+# Arguments
+max_t: maximum (minimum) terminal time for forward (backward) integration,
+respectively.
+"""
 mutable struct Clock
     t::Float64
     dt::Float64
@@ -51,8 +58,8 @@ function init!(clock::Clock, dt, ode_step_len; max_t=10.0)
     return clock
 end
 
-function reset!(clock::Clock)
-    clock.t = 0.0
+function reset!(clock::Clock; t=0.0)
+    clock.t = t
     return clock
 end
 
@@ -72,16 +79,17 @@ end
 "Check if the time is larger than max_t."
 function time_over(clock::Clock; t=nothing)
     if t == nothing
-        return time(clock) >= clock.max_t
+        return sign(clock.dt) * (time(clock) - clock.max_t) >= 0.0
     else
-        return t >= clock.max_t
+        return sign(clock.dt) * (t - clock.max_t) >= 0.0
     end
 end
 
 function thist(clock::Clock)
     thist = clock.thist .+ time(clock)
     if time_over(clock, t=thist[end])
-        index = findfirst(thist .> clock.max_t)
+        index = findfirst(sign(clock.dt)*(thist .- clock.max_t) .> 0.0)
+        # index = findfirst(thist .> clock.max_t)
         if index == nothing
             return thist
         else
@@ -106,21 +114,9 @@ end
 
 function _show(sys::BaseSystem; i=0)
     result = []
-    for symbol in [:name, :state, :dot, :initial_state,
-                   :state_size, :flat_index]
-        if isdefined(sys, symbol)
-            value = getproperty(sys, symbol)
-        else
-            value = "undef"
-        end
-        if symbol == :name
-            space = "+---"
-        else
-            space = "|   "
-        end
-        push!(result,
-              _add_space("$(String(symbol)): $value", i, space=space))
-    end
+    _stack_property!(result, sys,
+                     [:name, :state, :dot, :initial_state,
+                      :state_size, :flat_index], i=i)
     return join(result, "\n")
 end
 
@@ -169,6 +165,7 @@ mutable struct BaseEnv
     dyn
     step
 
+    initial_time
     dt::Float64
     clock::Clock
     progressbar
@@ -191,9 +188,29 @@ function _add_space(string, i; space=" "^4)
     return space^i * string
 end
 
+function _stack_property!(result, sys, symbol_array; i=0, space="")
+    for symbol in symbol_array
+        if isdefined(sys, symbol)
+            value = getproperty(sys, symbol)
+        else
+            value = "undef"
+        end
+        if symbol == :name
+            space = "+---"
+        else
+            space = "|   "
+        end
+        push!(result, _add_space("$(String(symbol)): $value", i, space=space))
+        # push!(result, _add_space("$(String(symbol)): $value", 0))
+    end
+end
+
 function _show(env::BaseEnv; i=0)
     result = []
-    push!(result, _add_space("name: $(env.name)", i, space="+---"))
+    _stack_property!(result, env, [:name])  # from env
+    if i == 0
+        _stack_property!(result, env.clock, [:max_t, :dt])  # from env.clock
+    end
     for system in _systems(env)
         if typeof(system) == BaseSystem
             v_str = _show(system, i=i+1)
@@ -213,15 +230,16 @@ end
 function init!(env::BaseEnv;
                systems=Dict(), dyn=nothing, step=nothing,
                # params=Dict(),
-               dt=0.01, max_t=1.0, ode_step_len=1,
+               initial_time=0.0, dt=0.01, max_t=1.0,
+               ode_step_len=1,
                logger=nothing, ode_option=Dict(), solver="rk4",
                name=nothing,
               )
     env.name = name
     systems!(env, systems)
 
     env.clock = Clock(dt, ode_step_len, max_t=max_t)
-    env.dt = env.clock.dt
+    env.initial_time = initial_time
 
     env.logger = logger
 
@@ -279,7 +297,7 @@ function reset!(env::BaseEnv)
     for system in _systems(env)
         reset!(system)
     end
-    reset!(env.clock)
+    reset!(env.clock, t=env.initial_time)
 end
 
 function state(env::BaseEnv)

test/runtests.jl

@@ -46,7 +46,7 @@ function test_Fym()
     log_dir = "data/test"
     file_name = "fym.h5"
     logger = Logger(log_dir=log_dir, file_name=file_name)
-    env = BaseEnv(max_t=100.00, logger=logger, name="test_env")
+    env = BaseEnv(max_t=100.00, logger=logger, name="test_env",)
     systems!(env, systems)  # set systems; required
     dyn!(env, set_dyn)  # set dynamics; required
     step!(env, step)  # set step; required
@@ -65,8 +65,61 @@ function test_Fym()
     end
     close!(env)
     data = load(env.logger.path)
-    @show env
-    @show size(data["state"]["sys"])
+    show(env)
+    show(size(data["state"]["sys"]))
+end
+
+function test_reverse_time()
+    print_msg("reverse_time")
+    function set_dyn(env, t)
+        # corresponding to `set_dot` of the original fym
+        # you can use any names in this package
+        sys = env.systems["sys"]
+        x = sys.state
+        A = Matrix(I, 3, 3)
+        sys.dot = -A * x
+    end
+    function step(env)
+        t = time(env.clock)
+        sys = env.systems["sys"]
+        x = sys.state
+        update!(env)
+        next_obs = sys.state
+        reward = zeros(1)
+        done = time_over(env.clock)
+        info = Dict(
+                    "time" => t,
+                    "state" => x,
+                   )
+        return next_obs, reward, done, info
+    end
+
+    x0 = collect(1:3)
+    systems = Dict("sys" => BaseSystem(initial_state=x0, name="3d_sys"))
+    log_dir = "data/test"
+    file_name = "fym.h5"
+    logger = Logger(log_dir=log_dir, file_name=file_name)
+    env = BaseEnv(max_t=0.00, dt=-0.01, initial_time=1.0, logger=logger, name="test_env",)
+    systems!(env, systems)  # set systems; required
+    dyn!(env, set_dyn)  # set dynamics; required
+    step!(env, step)  # set step; required
+
+    reset!(env)  # reset env; required before propagation
+    obs = observe_flat(env)
+    i = 0
+    @time while true
+        render(env)  # not mendatory; would make simulator slow
+        next_obs, reward, done, info = env.step()
+        obs = next_obs
+        i += 1
+        if done
+            break
+        end
+    end
+    close!(env)
+    data = load(env.logger.path)
+    show(env)
+    show(size(data["state"]["sys"]))
 end
 
 function test_largescale_env()
@@ -110,7 +163,7 @@ function test_largescale_env()
                    "sys" => BaseSystem(initial_state=x0),
                    "sys2" => BaseSystem(initial_state=y0),
                   )
-    env0 = BaseEnv(systems=systems0, dyn=set_dyn0)
+    env0 = BaseEnv(systems=systems0, dyn=set_dyn0, name="env0")
     systems = Dict(
                    "sys" => BaseSystem(initial_state=x0),
                    "sys2" => BaseSystem(initial_state=y0),
@@ -119,7 +172,7 @@ function test_largescale_env()
     log_dir = "data/test"
     file_name = "largescale.h5"
     logger = Logger(log_dir=log_dir, file_name=file_name)
-    env = BaseEnv(max_t=100.00, logger=logger)
+    env = BaseEnv(max_t=10.00, logger=logger, name="env")
     systems!(env, systems)
     dyn!(env, set_dyn)
     step!(env, step)
@@ -151,6 +204,7 @@ function test_largescale_env()
     savefig(p2, joinpath(log_dir, "state2.pdf"))
     savefig(env0p, joinpath(log_dir, "env0state.pdf"))
     savefig(env0p2, joinpath(log_dir, "env0state2.pdf"))
+    show(env)
 end
 
 function test_custom_env()
@@ -227,6 +281,7 @@ end
 
 function test_all()
     test_Fym()
+    test_reverse_time()
     test_largescale_env()
     test_custom_fym()
 end