Merge pull request #322 from SciML/fm/show

feat: add nice show for layers and models

Author: MartinuzziFrancesco

Project.toml

@@ -1,7 +1,7 @@
 name = "ReservoirComputing"
 uuid = "7c2d2b1e-3dd4-11ea-355a-8f6a8116e294"
 authors = ["Francesco Martinuzzi"]
-version = "0.12.2"
+version = "0.12.3"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/layers/basic.jl

@@ -338,6 +338,32 @@ function (dl::DelayLayer)(inp::AbstractVecOrMat, ps, st::NamedTuple)
     return inp_with_delay, (history = history, clock = clock, rng = st.rng)
 end
 
+function Base.show(io::IO, dl::DelayLayer)
+    print(io, "DelayLayer(", dl.in_dims, "; num_delays=", dl.num_delays)
+
+    if dl.stride != 1
+        print(io, ", stride=", dl.stride)
+    end
+
+    inits = dl.init_delay
+    if !isempty(inits)
+        all_same = all(f -> f === first(inits), inits)
+        if all_same && first(inits) === zeros32
+        elseif all_same
+            print(io, ", init_delay=", first(inits))
+        else
+            print(io, ", init_delay=(")
+            for (i, f) in enumerate(inits)
+                i > 1 && print(io, ", ")
+                show(io, f)
+            end
+            print(io, ")")
+        end
+    end
+
+    print(io, ")")
+end
+
 @doc raw"""
     NonlinearFeaturesLayer(features...; include_input=true)
 
@@ -417,3 +443,27 @@ function (nfl::NonlinearFeaturesLayer)(inp::AbstractVector, ps, st)
 
     return out, st
 end
+
+function Base.show(io::IO, nfl::NonlinearFeaturesLayer)
+    print(io, "NonlinearFeaturesLayer(")
+
+    if isempty(nfl.features)
+        print(io, "features=()")
+    else
+        print(io, "features=(")
+        for (i, f) in enumerate(nfl.features)
+            i > 1 && print(io, ", ")
+            show(io, f)
+        end
+        print(io, ")")
+    end
+
+    inc = ReservoirComputing.known(nfl.include_input)
+    if inc === true
+        print(io, ", include_input=true")
+    elseif inc === false
+        print(io, ", include_input=false")
+    end
+
+    print(io, ")")
+end

src/layers/esn_cell.jl

@@ -81,7 +81,7 @@ end
 function ESNCell((in_dims, out_dims)::Pair{<:IntegerType, <:IntegerType},
         activation = tanh; use_bias::BoolType = False(), init_bias = zeros32,
         init_reservoir = rand_sparse, init_input = scaled_rand,
-        init_state = randn32, leak_coefficient = 1.0)
+        init_state = randn32, leak_coefficient::AbstractFloat = 1.0)
     return ESNCell(activation, in_dims, out_dims, init_bias, init_reservoir,
         init_input, init_state, leak_coefficient, use_bias)
 end

src/layers/lux_layers.jl

@@ -45,6 +45,12 @@ function applyrecurrentcell(sl::AbstractReservoirRecurrentCell, inp, ps, st, ::N
     return apply(sl, inp, ps, st)
 end
 
+function Base.show(io::IO, sl::StatefulLayer)
+    print(io, "StatefulLayer(")
+    show(io, sl.cell)
+    print(io, ")")
+end
+
 @doc raw"""
     ReservoirChain(layers...; name=nothing)
     ReservoirChain(xs::AbstractVector; name=nothing)

src/models/esn.jl

@@ -102,5 +102,32 @@ function ESN(in_dims::IntegerType, res_dims::IntegerType,
                  Tuple(state_modifiers) : (state_modifiers,)
     mods = _wrap_layers(mods_tuple)
     ro = LinearReadout(res_dims => out_dims, readout_activation)
-    return ReservoirComputer(cell, mods, ro)
+    return ESN(cell, mods, ro)
+end
+
+function Base.show(io::IO, esn::ESN)
+    print(io, "ESN(\n")
+
+    print(io, "    reservoir = ")
+    show(io, esn.reservoir)
+    print(io, ",\n")
+
+    print(io, "    state_modifiers = ")
+    if isempty(esn.states_modifiers)
+        print(io, "()")
+    else
+        print(io, "(")
+        for (i, m) in enumerate(esn.states_modifiers)
+            i > 1 && print(io, ", ")
+            show(io, m)
+        end
+        print(io, ")")
+    end
+    print(io, ",\n")
+
+    print(io, "    readout = ")
+    show(io, esn.readout)
+    print(io, "\n)")
+
+    return
 end

src/models/esn_delay.jl

@@ -111,5 +111,32 @@ function DelayESN(
     ro_in_dims = res_dims * (num_delays + 1)
     ro = LinearReadout(ro_in_dims => out_dims, readout_activation)
 
-    return ReservoirComputer(cell, mods, ro)
+    return DelayESN(cell, mods, ro)
+end
+
+function Base.show(io::IO, esn::DelayESN)
+    print(io, "DelayESN(\n")
+
+    print(io, "    reservoir = ")
+    show(io, esn.reservoir)
+    print(io, ",\n")
+
+    print(io, "    state_modifiers = ")
+    if isempty(esn.states_modifiers)
+        print(io, "()")
+    else
+        print(io, "(")
+        for (i, m) in enumerate(esn.states_modifiers)
+            i > 1 && print(io, ", ")
+            show(io, m)
+        end
+        print(io, ")")
+    end
+    print(io, ",\n")
+
+    print(io, "    readout = ")
+    show(io, esn.readout)
+    print(io, "\n)")
+
+    return
 end

src/models/esn_hybrid.jl

@@ -56,7 +56,7 @@ throughout the reservoir and readout computations.
 ## Parameters
 
 - `knowledge_model` — parameters of the knowledge model `km`.
-- `cell` — parameters of the internal [`ESNCell`](@ref), including:
+- `reservoir` — parameters of the internal [`ESNCell`](@ref), including:
     - `input_matrix :: (res_dims × (in_dims + km_dims))` — `W_in`
     - `reservoir_matrix :: (res_dims × res_dims)` — `W_res`
     - `bias :: (res_dims,)` — present only if `use_bias=true`
@@ -73,13 +73,13 @@ throughout the reservoir and readout computations.
 Created by `initialstates(rng, hesn)`:
 
 - `knowledge_model` — states for the internal knowledge model.
-- `cell` — states for the internal [`ESNCell`](@ref).
+- `reservoir` — states for the internal [`ESNCell`](@ref).
 - `states_modifiers` — a `Tuple` with states for each modifier layer.
 - `readout` — states for [`LinearReadout`](@ref).
 """
 @concrete struct HybridESN <: AbstractEchoStateNetwork{(
-    :cell, :states_modifiers, :readout, :knowledge_model)}
-    cell
+    :reservoir, :states_modifiers, :readout, :knowledge_model)}
+    reservoir
     knowledge_model
     states_modifiers
     readout
@@ -94,46 +94,78 @@ function HybridESN(km,
         include_collect::BoolType = True(),
         kwargs...)
     esn_inp_size = in_dims + km_dims
-    cell = StatefulLayer(ESNCell(esn_inp_size => res_dims, activation; kwargs...))
+    reservoir = StatefulLayer(ESNCell(esn_inp_size => res_dims, activation; kwargs...))
     mods_tuple = state_modifiers isa Tuple || state_modifiers isa AbstractVector ?
                  Tuple(state_modifiers) : (state_modifiers,)
     mods = _wrap_layers(mods_tuple)
     ro = LinearReadout(res_dims + km_dims => out_dims, readout_activation;
         include_collect = static(include_collect))
     km_layer = km isa WrappedFunction ? km : WrappedFunction(km)
-    return HybridESN(cell, km_layer, mods, ro)
+    return HybridESN(reservoir, km_layer, mods, ro)
 end
 
 function initialparameters(rng::AbstractRNG, hesn::HybridESN)
-    ps_cell = initialparameters(rng, hesn.cell)
+    ps_reservoir = initialparameters(rng, hesn.reservoir)
     ps_km = initialparameters(rng, hesn.knowledge_model)
     ps_mods = map(l -> initialparameters(rng, l), hesn.states_modifiers) |> Tuple
     ps_ro = initialparameters(rng, hesn.readout)
-    return (cell = ps_cell, knowledge_model = ps_km,
+    return (reservoir = ps_reservoir, knowledge_model = ps_km,
         states_modifiers = ps_mods, readout = ps_ro)
 end
 
 function initialstates(rng::AbstractRNG, hesn::HybridESN)
-    st_cell = initialstates(rng, hesn.cell)
+    st_reservoir = initialstates(rng, hesn.reservoir)
     st_km = initialstates(rng, hesn.knowledge_model)
     st_mods = map(l -> initialstates(rng, l), hesn.states_modifiers) |> Tuple
     st_ro = initialstates(rng, hesn.readout)
-    return (cell = st_cell, knowledge_model = st_km,
+    return (reservoir = st_reservoir, knowledge_model = st_km,
         states_modifiers = st_mods, readout = st_ro)
 end
 
 function _partial_apply(hesn::HybridESN, inp, ps, st)
     k_t, st_km = hesn.knowledge_model(inp, ps.knowledge_model, st.knowledge_model)
     xin = vcat(k_t, inp)
-    r, st_cell = apply(hesn.cell, xin, ps.cell, st.cell)
+    r, st_reservoir = apply(hesn.reservoir, xin, ps.reservoir, st.reservoir)
     rstar,
     st_mods = _apply_seq(hesn.states_modifiers, r, ps.states_modifiers, st.states_modifiers)
     feats = vcat(k_t, rstar)
-    return feats, (cell = st_cell, states_modifiers = st_mods, knowledge_model = st_km)
+    return feats,
+    (reservoir = st_reservoir, states_modifiers = st_mods, knowledge_model = st_km)
 end
 
 function (hesn::HybridESN)(inp, ps, st)
     feats, new_st = _partial_apply(hesn, inp, ps, st)
     y, st_ro = apply(hesn.readout, feats, ps.readout, st.readout)
     return y, merge(new_st, (readout = st_ro,))
 end
+
+function Base.show(io::IO, esn::HybridESN)
+    print(io, "HybridESN(\n")
+
+    print(io, "    reservoir = ")
+    show(io, esn.reservoir)
+    print(io, ",\n")
+
+    print(io, "    knowledge_model = ")
+    show(io, esn.knowledge_model)
+    print(io, ",\n")
+
+    print(io, "    state_modifiers = ")
+    if isempty(esn.states_modifiers)
+        print(io, "()")
+    else
+        print(io, "(")
+        for (i, m) in enumerate(esn.states_modifiers)
+            i > 1 && print(io, ", ")
+            show(io, m)
+        end
+        print(io, ")")
+    end
+    print(io, ",\n")
+
+    print(io, "    readout = ")
+    show(io, esn.readout)
+    print(io, "\n)")
+
+    return
+end

src/reservoircomputer.jl

@@ -97,6 +97,30 @@ function addreadout!(::AbstractReservoirComputer, output_matrix::AbstractMatrix,
     return merge(ps, (readout = new_readout,)), st
 end
 
+function Base.show(io::IO, rc::ReservoirComputer)
+    print(io, "ReservoirComputer(")
+
+    print(io, "reservoir = ")
+    show(io, rc.reservoir)
+
+    nmods = length(rc.states_modifiers)
+    if nmods == 0
+        print(io, ", state_modifiers = ()")
+    else
+        print(io, ", state_modifiers = (")
+        for (i, m) in enumerate(rc.states_modifiers)
+            i > 1 && print(io, ", ")
+            show(io, m)
+        end
+        print(io, ")")
+    end
+
+    print(io, ", readout = ")
+    show(io, rc.readout)
+
+    print(io, ")")
+end
+
 @doc raw"""
     resetcarry!(rng, rc::ReservoirComputer, st; init_carry=nothing)
     resetcarry!(rng, rc::ReservoirComputer, ps, st; init_carry=nothing)