Merge pull request #136 from SciML/fm/states_pred

Adding `save_states`

Project.toml

@@ -1,7 +1,7 @@
 name = "ReservoirComputing"
 uuid = "7c2d2b1e-3dd4-11ea-355a-8f6a8116e294"
 authors = ["Francesco Martinuzzi"]
-version = "0.9.0"
+version = "0.9.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/esn_tutorials/lorenz_basic.md

@@ -86,6 +86,10 @@ Once the ```OutputLayer``` has been obtained the prediction can be done followin
 output = esn(Generative(predict_len), output_layer)
 ```
 both the training method and the output layer are needed in this call. The number of steps for the prediction must be specified to the ```Generative``` method. The output results are given in a matrix. 
+
+!!! info "Saving the states during prediction"
+    While the states are saved in the `ESN` struct for the training, for the prediction they are not saved by default. To inspect the states it is necessary to pass the boolean keyword argument `save_states` to the prediction call, in this example using `esn(... ; save_states=true)`. This returns a tuple `(output, states)` where `size(states) = res_size, prediction_len`
+
 To inspect the results they can easily be plotted using an external library. In this case ```Plots``` is adopted:
 ```julia
 using Plots, Plots.PlotMeasures

src/esn/echostatenetwork.jl

@@ -164,8 +164,8 @@ end
 
 function (esn::ESN)(prediction::AbstractPrediction,
                     output_layer::AbstractOutputLayer;
-                    initial_conditions = output_layer.last_value,
-                    last_state = esn.states[:, [end]])
+                    last_state = esn.states[:, [end]],
+                    kwargs...)
     variation = esn.variation
     pred_len = prediction.prediction_len
 
@@ -178,10 +178,10 @@ function (esn::ESN)(prediction::AbstractPrediction,
         model_pred_data = model(u0, tspan_new, predict_tsteps)[:, 2:end]
         return obtain_esn_prediction(esn, prediction, last_state, output_layer,
                                      model_pred_data;
-                                     initial_conditions = initial_conditions)
+                                     kwargs...)
     else
         return obtain_esn_prediction(esn, prediction, last_state, output_layer;
-                                     initial_conditions = initial_conditions)
+                                     kwargs...)
     end
 end
 

src/esn/esn_predict.jl

@@ -3,13 +3,15 @@ function obtain_esn_prediction(esn,
                                x,
                                output_layer,
                                args...;
-                               initial_conditions = output_layer.last_value)
+                               initial_conditions = output_layer.last_value,
+                               save_states = false)
     out_size = output_layer.out_size
     training_method = output_layer.training_method
     prediction_len = prediction.prediction_len
 
     output = output_storing(training_method, out_size, prediction_len, typeof(esn.states))
     out = initial_conditions
+    states = similar(esn.states, size(esn.states, 1), prediction_len)
 
     out_pad = allocate_outpad(esn.variation, esn.states_type, out)
     tmp_array = allocate_tmp(esn.reservoir_driver, typeof(esn.states), esn.res_size)
@@ -20,23 +22,26 @@ function obtain_esn_prediction(esn,
                                           args...)
         out_tmp = get_prediction(output_layer.training_method, output_layer, x_new)
         out = store_results!(output_layer.training_method, out_tmp, output, i)
+        states[:, i] = x
     end
 
-    return output
+    save_states ? (output, states) : output
 end
 
 function obtain_esn_prediction(esn,
                                prediction::Predictive,
                                x,
                                output_layer,
                                args...;
-                               initial_conditions = output_layer.last_value)
+                               initial_conditions = output_layer.last_value,
+                               save_states = false)
     out_size = output_layer.out_size
     training_method = output_layer.training_method
     prediction_len = prediction.prediction_len
 
     output = output_storing(training_method, out_size, prediction_len, typeof(esn.states))
     out = initial_conditions
+    states = similar(esn.states, size(esn.states, 1), prediction_len)
 
     out_pad = allocate_outpad(esn.variation, esn.states_type, out)
     tmp_array = allocate_tmp(esn.reservoir_driver, typeof(esn.states), esn.res_size)
@@ -47,9 +52,10 @@ function obtain_esn_prediction(esn,
                                           out_pad, i, tmp_array, args...)
         out_tmp = get_prediction(training_method, output_layer, x_new)
         out = store_results!(training_method, out_tmp, output, i)
+        states[:, i] = x
     end
 
-    return output
+    save_states ? (output, states) : output
 end
 
 #prediction dispatch on esn 

test/esn/test_train.jl

@@ -26,3 +26,9 @@ for t in training_methods
     output = esn(Predictive(input_data), output_layer)
     @test mean(abs.(target_data .- output)) ./ mean(abs.(target_data)) < 0.21
 end
+
+for t in training_methods
+    output_layer = train(esn, target_data, t)
+    output, states = esn(Predictive(input_data), output_layer, save_states = true)
+    @test size(states) == (res_size, size(input_data, 2))
+end