refac: build on LuxCore

Project.toml

@@ -1,15 +1,18 @@
 name = "ReservoirComputing"
 uuid = "7c2d2b1e-3dd4-11ea-355a-8f6a8116e294"
 authors = ["Francesco Martinuzzi"]
-version = "0.11.4"
+version = "0.12.0"
 
 [deps]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 
 [weakdeps]
@@ -25,19 +28,23 @@ RCMLJLinearModelsExt = "MLJLinearModels"
 RCSparseArraysExt = "SparseArrays"
 
 [compat]
-Adapt = "4.1.1"
 Aqua = "0.8"
+ArrayInterface = "7.19.0"
 CellularAutomata = "0.0.6"
 Compat = "4.16.0"
+ConcreteStructs = "0.2.3"
 DifferentialEquations = "7.16.1"
+JET = "0.9, 0.10.10"
 LIBSVM = "0.8"
 LinearAlgebra = "1.10"
+LuxCore = "1.3.0"
 MLJLinearModels = "0.9.2, 0.10"
 NNlib = "0.9.26"
 Random = "1.10"
 Reexport = "1.2.2"
 SafeTestsets = "0.1"
 SparseArrays = "1.10"
+Static = "1.2.0"
 Statistics = "1.10"
 Test = "1"
 WeightInitializers = "1.0.5"
@@ -46,6 +53,7 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LIBSVM = "b1bec4e5-fd48-53fe-b0cb-9723c09d164b"
 MLJLinearModels = "6ee0df7b-362f-4a72-a706-9e79364fb692"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -54,4 +62,4 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "SafeTestsets", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "Statistics", "SparseArrays", "CellularAutomata"]
+test = ["Aqua", "Test", "SafeTestsets", "DifferentialEquations", "MLJLinearModels", "LIBSVM", "Statistics", "SparseArrays", "JET"]

README.md

@@ -27,22 +27,18 @@ Use the
 [in-development documentation](https://docs.sciml.ai/ReservoirComputing/dev/)
 to take a look at not yet released features.
 
-## Citing
+## Features
 
-If you use this library in your work, please cite:
+ReservoirComputing.jl provides layers,models, and functions to help build and train
+reservoir computing models. More specifically the software offers
 
-```bibtex
-@article{martinuzzi2022reservoircomputing,
-  author  = {Francesco Martinuzzi and Chris Rackauckas and Anas Abdelrehim and Miguel D. Mahecha and Karin Mora},
-  title   = {ReservoirComputing.jl: An Efficient and Modular Library for Reservoir Computing Models},
-  journal = {Journal of Machine Learning Research},
-  year    = {2022},
-  volume  = {23},
-  number  = {288},
-  pages   = {1--8},
-  url     = {http://jmlr.org/papers/v23/22-0611.html}
-}
-```
+- Base layers for reservoir computing model construction such as `ReservoirChain`,
+  `Readout`, `Collect`, and `ESNCell`
+- Fully built models such as `ESN`, and `DeepESN`
+- 15+ reservoir initializers and 5+ input layer initializers
+- 5+ reservoir states modification algorithms
+- Sparse matrix computation through
+  [SparseArrays.jl](https://docs.julialang.org/en/v1/stdlib/SparseArrays/)
 
 ## Installation
 
@@ -63,70 +59,82 @@ Pkg.add("ReservoirComputing")
 
 To illustrate the workflow of this library we will showcase
 how it is possible to train an ESN to learn the dynamics of the
-Lorenz system. As a first step we gather the data.
-For the `Generative` prediction we need the target data
-to be one step ahead of the training data:
+Lorenz system.
+
+### 1. Generate data
+
+As a general first step wee fix the random seed for reproducibilty
 
 ```julia
-using ReservoirComputing, OrdinaryDiffEq, Random
+using Random
 Random.seed!(42)
 rng = MersenneTwister(17)
+```
 
-#lorenz system parameters
-u0 = [1.0, 0.0, 0.0]
-tspan = (0.0, 200.0)
-p = [10.0, 28.0, 8 / 3]
+For an autoregressive prediction we need the target data
+to be one step ahead of the training data:
+
+```julia
+using OrdinaryDiffEq
 
 #define lorenz system
 function lorenz(du, u, p, t)
     du[1] = p[1] * (u[2] - u[1])
     du[2] = u[1] * (p[2] - u[3]) - u[2]
     du[3] = u[1] * u[2] - p[3] * u[3]
 end
+
 #solve and take data
-prob = ODEProblem(lorenz, u0, tspan, p)
+prob = ODEProblem(lorenz, [1.0f0, 0.0f0, 0.0f0], (0.0, 200.0), [10.0f0, 28.0f0, 8/3])
 data = Array(solve(prob, ABM54(); dt=0.02))
-
 shift = 300
 train_len = 5000
 predict_len = 1250
 
 #one step ahead for generative prediction
 input_data = data[:, shift:(shift + train_len - 1)]
 target_data = data[:, (shift + 1):(shift + train_len)]
-
 test = data[:, (shift + train_len):(shift + train_len + predict_len - 1)]
 ```
 
-Now that we have the data we can initialize the ESN with the chosen parameters.
-Given that this is a quick example we are going to change the least amount of
-possible parameters:
+### 2. Build Echo State Network
+
+We can either use the provided `ESN` or build one from scratch.
+We showcase the second option:
 
 ```julia
+using ReservoirComputing
 input_size = 3
 res_size = 300
-esn = ESN(input_data, input_size, res_size;
-    reservoir=rand_sparse(; radius=1.2, sparsity=6 / res_size),
-    input_layer=weighted_init,
-    nla_type=NLAT2(),
-    rng=rng)
+esn = ESN(input_size, res_size, input_size;
+    init_reservoir=rand_sparse(; radius=1.2, sparsity=6/300),
+    state_modifiers=NLAT2
+)
+```
+
+### 3. Train the Echo State Network
+
+ReservoirComputing.jl builds on Lux(Core), so in order to train the model
+we first need to instantiate the parameters and the states:
+
+```julia
+ps, st = setup(rng, esn)
+ps, st = train!(esn, input_data, target_data, ps, st)
 ```
 
-The echo state network can now be trained and tested.
-If not specified, the training will always be ordinary least squares regression:
+### 4. Predict and visualize
+
+We can now use the trained ESN to forecast the Lorenz system dynamics
 
 ```julia
-output_layer = train(esn, target_data)
-output = esn(Generative(predict_len), output_layer)
+output, st = predict(esn, 1250, ps, st; initialdata=test[:, 1])
 ```
 
-The data is returned as a matrix, `output` in the code above,
-that contains the predicted trajectories.
-The results can now be easily plotted:
+We can now visualize the results
 
 ```julia
 using Plots
-plot(transpose(output); layout=(3, 1), label="predicted")
+plot(transpose(output); layout=(3, 1), label="predicted");
 plot!(transpose(test); layout=(3, 1), label="actual")
 ```
 
@@ -145,6 +153,23 @@ plot!(transpose(test)[:, 1], transpose(test)[:, 2], transpose(test)[:, 3]; label
 
 ![lorenz_attractor](https://user-images.githubusercontent.com/10376688/81470281-5a34b580-91ea-11ea-9eea-d2b266da19f4.png)
 
+## Citing
+
+If you use this library in your work, please cite:
+
+```bibtex
+@article{martinuzzi2022reservoircomputing,
+  author  = {Francesco Martinuzzi and Chris Rackauckas and Anas Abdelrehim and Miguel D. Mahecha and Karin Mora},
+  title   = {ReservoirComputing.jl: An Efficient and Modular Library for Reservoir Computing Models},
+  journal = {Journal of Machine Learning Research},
+  year    = {2022},
+  volume  = {23},
+  number  = {288},
+  pages   = {1--8},
+  url     = {http://jmlr.org/papers/v23/22-0611.html}
+}
+```
+
 ## Acknowledgements
 
 This project was possible thanks to initial funding through

docs/Project.toml

@@ -1,23 +1,20 @@
 [deps]
 CellularAutomata = "878138dc-5b27-11ea-1a71-cb95d38d6b29"
-DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
+DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
-PredefinedDynamicalSystems = "31e2f376-db9e-427a-b76e-a14f56347a14"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReservoirComputing = "7c2d2b1e-3dd4-11ea-355a-8f6a8116e294"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 
 [compat]
 CellularAutomata = "0.0.6"
-DifferentialEquations = "7.16.1"
 Documenter = "1"
 DocumenterCitations = "1"
 OrdinaryDiffEq = "6"
 Plots = "1"
-PredefinedDynamicalSystems = "1"
-ReservoirComputing = "0.11.4"
+ReservoirComputing = "0.12.0"
 StatsBase = "0.34.4"

docs/make.jl

@@ -1,7 +1,7 @@
-using Documenter, DocumenterCitations, ReservoirComputing
+using Documenter, DocumenterCitations, DocumenterInterLinks, ReservoirComputing
 
-cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml"; force = true)
-cp("./docs/Project.toml", "./docs/src/assets/Project.toml"; force = true)
+#cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml"; force=true)
+#cp("./docs/Project.toml", "./docs/src/assets/Project.toml"; force=true)
 
 ENV["PLOTS_TEST"] = "true"
 ENV["GKSwstype"] = "100"
@@ -13,10 +13,14 @@ bib = CitationBibliography(
     style = :authoryear
 )
 
+links = InterLinks(
+    "Lux" => "https://lux.csail.mit.edu/stable/"
+)
+
 makedocs(; modules = [ReservoirComputing],
     sitename = "ReservoirComputing.jl",
     clean = true, doctest = false, linkcheck = true,
-    plugins = [bib],
+    plugins = [links, bib],
     format = Documenter.HTML(;
         mathengine,
         assets = ["assets/favicon.ico"],

docs/pages.jl

@@ -1,24 +1,19 @@
 pages = [
     "ReservoirComputing.jl" => "index.md",
-    "General Settings" => Any[
-        "Changing Training Algorithms" => "general/different_training.md",
-        "Altering States" => "general/states_variation.md",
-        "Generative vs Predictive" => "general/predictive_generative.md"],
-    "Echo State Network Tutorials" => Any[
-        "Lorenz System Forecasting" => "esn_tutorials/lorenz_basic.md",
-        #"Mackey-Glass Forecasting on GPU" => "esn_tutorials/mackeyglass_basic.md",
-        "Using Different Layers" => "esn_tutorials/change_layers.md",
-        "Using Different Reservoir Drivers" => "esn_tutorials/different_drivers.md",
+    "Tutorials" => Any[
+        "Building a model from scratch" => "tutorials/scratch.md",
+        "Chaos forecasting with an ESN" => "tutorials/lorenz_basic.md",
         #"Using Different Training Methods" => "esn_tutorials/different_training.md",
-        "Deep Echo State Networks" => "esn_tutorials/deep_esn.md",
-        "Hybrid Echo State Networks" => "esn_tutorials/hybrid.md"],
-    "Reservoir Computing with Cellular Automata" => "reca_tutorials/reca.md",
-    "API Documentation" => Any["Training Algorithms" => "api/training.md",
-        "States Modifications" => "api/states.md",
-        "Prediction Types" => "api/predict.md",
-        "Echo State Networks" => "api/esn.md",
-        "ESN Initializers" => "api/inits.md",
-        "ESN Drivers" => "api/esn_drivers.md",
-        "ESN Variations" => "api/esn_variations.md",
-        "ReCA" => "api/reca.md"]    #"References" => "references.md"
+        "Deep Echo State Networks" => "tutorials/deep_esn.md",
+        #"Hybrid Echo State Networks" => "tutorials/hybrid.md",
+        "Reservoir Computing with Cellular Automata" => "tutorials/reca.md"],
+    "API Documentation" => Any[
+        "Layers" => "api/layers.md",
+        "Models" => "api/models.md",
+        "Utilities" => "api/utils.md",
+        "Train" => "api/train.md",
+        "Predict" => "api/predict.md",
+        "States" => "api/states.md",
+        "Initializers" => "api/inits.md"],
+    "References" => "references.md"
 ]