Inference of a DigitalEcho

Once we train a DigitalEcho, it is important to understand as to how to use it for predictions, i.e, inference or forward pass.

The way we do inference is slightly different for full model and component DigitalEcho. Let us go through both of them.

Full Model

First let us load a DigitalEcho trained on lotka-volterra data and also the data which we will use it for inference.

using JuliaHub, JLSO, Surrogatize, DataGeneration
digitalecho_dataset_name = "lotka_volterra_digitalecho"
path = JuliaHub.download_dataset(("juliasimtutorials", digitalecho_dataset_name), "path to save")
digitalecho = JLSO.load(path)[:result]
A Continuous Time Surrogate wrapper with:
prob:
  A `DigitalEchoProblem` with:
  model:
    A DigitalEcho with : 
      RSIZE : 256
      USIZE : 2
      XSIZE : 0
      PSIZE : 4
      ICSIZE : 0
solver: Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),)
train_dataset_name = "lotka_volterra"
path = JuliaHub.download_dataset(("juliasimtutorials", train_dataset_name), "path to save")
ed = ExperimentData(JLSO.load(path)[:result])
 Number of Trajectories in ExperimentData: 10 
  Basic Statistics for Given Dynamical System's Specifications 
  Number of u0s in the ExperimentData: 2 
  Number of ps in the ExperimentData: 4 
 ╭─────────┬──────────────────────────────────────────────────────────────────...
────╮...
  Field  ...
         ...
├─────────┼──────────────────────────────────────────────────────────────────...
────┤...
             ╭────────────┬──────────────┬──────────────┬────────┬─────────...
                Labels     LowerBound    UpperBound    Mean    StdDev...
             ├────────────┼──────────────┼──────────────┼────────┼─────────...
               states_1        1             1          1    ...
   u0s       ├────────────┼──────────────┼──────────────┼────────┼─────────...
             ...
             ...
             ├────────────┼──────────────┼──────────────┼────────┼─────────...
               states_2        1             1          1    ...
             ╰────────────┴──────────────┴──────────────┴────────┴─────────...
├─────────┼──────────────────────────────────────────────────────────────────...
────┤...
           ╭──────────┬──────────────┬──────────────┬───────────┬──────────...
             Labels    LowerBound    UpperBound     Mean       StdDev...
           ├──────────┼──────────────┼──────────────┼───────────┼──────────...
              p_1        1.5625        2.4375      1.96875    0.301904...
   ps      ├──────────┼──────────────┼──────────────┼───────────┼──────────...
           ...
           ...
           ├──────────┼──────────────┼──────────────┼───────────┼──────────...
              p_4       1.76562       1.98438      1.87031    0.074316...
           ╰──────────┴──────────────┴──────────────┴───────────┴──────────...
╰─────────┴──────────────────────────────────────────────────────────────────...
────╯...
 Basic Statistics for Given Dynamical System's Continuous Fields 
  Number of states in the ExperimentData: 2 
 ╭──────────┬─────────────────────────────────────────────────────────────────...
──────╮...
  Field   ...
           ...
├──────────┼─────────────────────────────────────────────────────────────────...
──────┤...
            ╭────────────┬──────────────┬──────────────┬───────────┬───────...
               Labels     LowerBound    UpperBound     Mean    ...
            ├────────────┼──────────────┼──────────────┼───────────┼───────...
              states_1     0.60988       1.85127      1.13137  ...
  states    ├────────────┼──────────────┼──────────────┼───────────┼───────...
            ...
            ...
            ├────────────┼──────────────┼──────────────┼───────────┼───────...
              states_2     0.585184      1.92984      1.0678   ...
            ╰────────────┴──────────────┴──────────────┴───────────┴───────...
╰──────────┴─────────────────────────────────────────────────────────────────...
──────╯...

To do inference, we need to index into the dataset appropriately and pass into the forward pass of the DigitalEcho. Let us try with the first sample of the dataset.

Firstly, let us get the initial condition.

u0 = ed.specs.u0s.vals[1]
2-element Vector{Float64}:
 1.0
 1.0

Next, we need to get the control function. As it is a full model, we don't have any controls. The function signature for passing control always takes in both state and time as arguments. Hence we create a closure and return nothing.

x = (u, t) -> nothing
#1 (generic function with 1 method)

Then, let us grab the parameters.

p = ed.specs.ps.vals[1]
4-element view(::Matrix{Float64}, :, 1) with eltype Float64:
 1.6875
 1.828125
 2.4375
 1.859375

Lastly, let us grab the timepoints we want to save the values and the time span.

ts = ed.results.tss.vals[1]
tspan = (ts[1], ts[end])
(0.0, 12.5)

We can then call the forward pass.

pred = digitalecho(u0, x, p, tspan; saveat = ts)
2×94 Matrix{Float64}:
 0.999981  0.994967  0.990948  …  1.09564  1.03981  1.00235  0.997773
 1.0       0.975376  0.938141     1.16734  1.09714  1.0091   0.990246

Component

First let us load a DigitalEcho trained on CSTR data and also the data which we will use it for inference.

using JuliaHub, JLSO, Surrogatize, DataGeneration, PreProcessing
digitalecho_dataset_name = "cstr_digitalecho"
path = JuliaHub.download_dataset(digitalecho_dataset, "path to save")
digitalecho = JLSO.load(path)[:result]
A Continuous Time Surrogate wrapper with:
prob:
  A `DigitalEchoProblem` with:
  model:
    A DigitalEcho with : 
      RSIZE : 256
      USIZE : 4
      XSIZE : 2
      PSIZE : 0
      ICSIZE : 0
solver: Tsit5(; stage_limiter! = trivial_limiter!, step_limiter! = trivial_limiter!, thread = static(false),)
train_dataset_name = "cstr"
path = JuliaHub.download_dataset(("juliasimtutorials", train_dataset_name), "path to save")
ed = ExperimentData(JLSO.load(path)[:result])
 Number of Trajectories in ExperimentData: 10 
  Basic Statistics for Given Dynamical System's Specifications 
  Number of u0s in the ExperimentData: 4 
 ╭─────────┬──────────────────────────────────────────────────────────────────...
─────╮...
  Field  ...
          ...
├─────────┼──────────────────────────────────────────────────────────────────...
─────┤...
           ╭────────────┬──────────────┬──────────────┬────────┬───────────...
              Labels     LowerBound    UpperBound    Mean      StdDev...
           ├────────────┼──────────────┼──────────────┼────────┼───────────...
             states_1       0.8           0.8        0.8     1.17028e-...
   u0s     ├────────────┼──────────────┼──────────────┼────────┼───────────...
           ...
           ...
           ├────────────┼──────────────┼──────────────┼────────┼───────────...
             states_4       130           130        130   ...
           ╰────────────┴──────────────┴──────────────┴────────┴───────────...
╰─────────┴──────────────────────────────────────────────────────────────────...
─────╯...
 Basic Statistics for Given Dynamical System's Continuous Fields 
  Number of states in the ExperimentData: 4 
  Number of controls in the ExperimentData: 2 
 ╭────────────┬───────────────────────────────────────────────────────────────...
────────╮...
   Field    ...
             ...
├────────────┼───────────────────────────────────────────────────────────────...
────────┤...
              ╭────────────┬──────────────┬──────────────┬───────────┬─────...
                 Labels     LowerBound    UpperBound     Mean    ...
              ├────────────┼──────────────┼──────────────┼───────────┼─────...
                states_1     0.782032       2.6942      2.02012  ...
   states     ├────────────┼──────────────┼──────────────┼───────────┼─────...
              ...
              ...
              ├────────────┼──────────────┼──────────────┼───────────┼─────...
                states_4     125.364       139.993      131.565  ...
              ╰────────────┴──────────────┴──────────────┴───────────┴─────...
├────────────┼───────────────────────────────────────────────────────────────...
────────┤...
               ╭──────────┬──────────────┬──────────────┬───────────┬──────...
                 Labels    LowerBound    UpperBound     Mean    ...
               ├──────────┼──────────────┼──────────────┼───────────┼──────...
                   F        8.33348       99.4681      55.8962  ...
  controls     ├──────────┼──────────────┼──────────────┼───────────┼──────...
               ...
               ...
               ├──────────┼──────────────┼──────────────┼───────────┼──────...
                  -8497.87      -227.416     -4158.5  ...
               ╰──────────┴──────────────┴──────────────┴───────────┴──────...
╰────────────┴───────────────────────────────────────────────────────────────...
────────╯...

Before doing inference, we need to convert the dataset into splines for continuous controls.

spline = SplineED()
ed_splines = spline(ed)
 Number of Trajectories in ExperimentData: 10 
  Basic Statistics for Given Dynamical System's Specifications 
  Number of u0s in the ExperimentData: 4 
 ╭─────────┬──────────────────────────────────────────────────────────────────...
─────╮...
  Field  ...
          ...
├─────────┼──────────────────────────────────────────────────────────────────...
─────┤...
           ╭────────────┬──────────────┬──────────────┬────────┬───────────...
              Labels     LowerBound    UpperBound    Mean      StdDev...
           ├────────────┼──────────────┼──────────────┼────────┼───────────...
             states_1       0.8           0.8        0.8     1.17028e-...
   u0s     ├────────────┼──────────────┼──────────────┼────────┼───────────...
           ...
           ...
           ├────────────┼──────────────┼──────────────┼────────┼───────────...
             states_4       130           130        130   ...
           ╰────────────┴──────────────┴──────────────┴────────┴───────────...
╰─────────┴──────────────────────────────────────────────────────────────────...
─────╯...
 Basic Statistics for Given Dynamical System's Continuous Fields 
  Number of states in the ExperimentData: 4 
  Number of controls in the ExperimentData: 2 
 ╭────────────┬───────────────────────────────────────────────────────────────...
────────╮...
   Field    ...
             ...
├────────────┼───────────────────────────────────────────────────────────────...
────────┤...
              ╭────────────┬──────────────┬──────────────┬───────────┬─────...
                 Labels     LowerBound    UpperBound     Mean    ...
              ├────────────┼──────────────┼──────────────┼───────────┼─────...
                states_1     0.782032       2.6942      2.02012  ...
   states     ├────────────┼──────────────┼──────────────┼───────────┼─────...
              ...
              ...
              ├────────────┼──────────────┼──────────────┼───────────┼─────...
                states_4     125.364       139.993      131.565  ...
              ╰────────────┴──────────────┴──────────────┴───────────┴─────...
├────────────┼───────────────────────────────────────────────────────────────...
────────┤...
               ╭──────────┬──────────────┬──────────────┬───────────┬──────...
                 Labels    LowerBound    UpperBound     Mean    ...
               ├──────────┼──────────────┼──────────────┼───────────┼──────...
                   F        8.33348       99.4681      55.8962  ...
  controls     ├──────────┼──────────────┼──────────────┼───────────┼──────...
               ...
               ...
               ├──────────┼──────────────┼──────────────┼───────────┼──────...
                  -8497.87      -227.416     -4158.5  ...
               ╰──────────┴──────────────┴──────────────┴───────────┴──────...
╰────────────┴───────────────────────────────────────────────────────────────...
────────╯...

Once we have converted into splines, next step is to index into various fields similar to how we do inference in the full model.

Firstly, let us get the initial condition.

u0 = ed_splines.specs.u0s.vals[1]
4-element Vector{Float64}:
   0.8
   0.5
 134.14
 130.0

Next, we get the controls. The dataset was generated using open loop controls. The function signature for passing control always takes in both state and time as arguments. Hence we create a closure with the spline we index from the dataset.

x = (u, t) -> ed_splines.results.controls.vals[1](t)
#3 (generic function with 1 method)

We set the parameters as nothing as the DigitalEcho was only trained on controls and cannot vary parameters.

p = nothing

Lastly, let us grab the timepoints we want to save the values and the time span.

ts = ed_splines.results.tss.vals[1]
tspan = (ts[1], ts[end])
(0.0, 0.25)

We can then call the forward pass.

pred = digitalecho(u0, x, p, tspan; saveat = ts)
4×48 Matrix{Float64}:
   0.799784    1.17452     1.34899   …    2.16619    2.16411    2.16226
   0.499969    0.497726    0.511656       1.11118    1.11157    1.11182
 134.14      133.659     133.463        135.057    135.083    135.107
 130.0       129.811     129.682        125.517    125.628    125.717