Plot option in batch.jl function batchDataSolution uses wrong index

In src/batch.jl in function batchDataSolution it is written:

        if plot
            fig = FMIFlux.plot(batch[i-1]; solverKwargs...)
            display(fig)
        end

instead it should be

        if plot
            fig = FMIFlux.plot(batch[i]; solverKwargs...)
            display(fig)
        end

If you call batchDataSolution with option plot=true, you immediately get an error, because there is no batch[0]. Furthermore, I don't see any reason why you want to plot any other batch than the one you just simulated.

Below is a reduced variant of the juliacon_2023.ipynb, I know it is not minimal, but the error is obvious, any more reduction of the script to generate the error is not worth the effort.

# Loading the required libraries
using FMI           # import FMUs into Julia 
using FMIFlux       # for NeuralFMUs
using FMIZoo        # a collection of demo models, including the VLDM
using FMIFlux.Flux  # Machine Learning in Julia

import JLD2         # data format for saving/loading parameters

import Random       # for fixing the random seed
using Plots         # plotting results

# a helper file with some predefined functions to make "things look nicer", but are not really relevant to the topic
include(joinpath(@__DIR__, "juliacon_2023_helpers.jl"));

showProgress=false;
dt = 0.1 
data = VLDM(:train, dt=dt)
############################################################

############################################################
#->PREPARATION OF SIMULATION DATA (for NN) ->manipulatedDerVals
# start (`tStart`) and stop time (`tStop`) for simulation, saving time points for ODE solver (`tSave`)
tStart = data.consumption_t[1]
tStop = data.consumption_t[end]
tSave = data.consumption_t
fmu = fmiLoad("VLDM", "Dymola", "2020x"; type=:ME) 
resultFMU = fmiSimulate(fmu,                        # the loaded FMU of the VLDM 
                        (tStart, tStop);            # the simulation time range
                        parameters=data.params,     # the parameters for the VLDM
                        showProgress=showProgress,  # show (or don't) the progres bar
                        recordValues=:derivatives,  # record all state derivatives
                        saveat=tSave)               # save solution points at `tSave`

# variable we want to manipulate - why we are picking exactly these three is shown a few lines later ;-)
manipulatedDerVars = ["der(dynamics.accelerationCalculation.integrator.y)",
                    "der(dynamics.accelerationCalculation.limIntegrator.y)",
                    "der(result.integrator.y)"]
manipulatedDerVals = fmiGetSolutionValue(resultFMU, manipulatedDerVars)
# unload FMU 
fmiUnload(fmu)
#############################################################################

###############################################################
# function that builds the considered NeuralFMU on basis of a given FMU (FMI-Version 2.0) `f`
function build_NFMU(f::FMU2)

    # pre- and post-processing
    preProcess = ShiftScale(manipulatedDerVals)         # we put in the derivatives recorded above, FMIFlux shift and scales so we have a data mean of 0 and a standard deivation of 1
    preProcess.scale[:] *= 0.25                         # add some additional "buffer"
    postProcess = ScaleShift(preProcess; indices=2:3)   # initialize the postPrcess as inverse of the preProcess, but only take indices 2 and 3 (we don't need 1, the vehcile velocity)

    # cache
    cache = CacheLayer()                        # allocate a cache layer
    cacheRetrieve = CacheRetrieveLayer(cache)   # allocate a cache retrieve layer, link it to the cache layer

    # we have two signals (acceleration, consumption) and two sources (ANN, FMU), so four gates:
    # (1) acceleration from FMU (gate=1.0 | open)
    # (2) consumption  from FMU (gate=1.0 | open)
    # (3) acceleration from ANN (gate=0.0 | closed)
    # (4) consumption  from ANN (gate=0.0 | closed)
    # the acelerations [1,3] and consumptions [2,4] are paired
    gates = ScaleSum([1.0, 1.0, 0.0, 0.0], [[1,3], [2,4]]) # gates with sum

    # setup the NeuralFMU topology
    model = Chain(x -> f(; x=x, dx_refs=:all),        # take `x`, put it into the FMU, retrieve all derivatives `dx`
                dx -> cache(dx),                    # cache `dx`
                dx -> dx[4:6],                      # forward only dx[4, 5, 6]
                preProcess,                         # pre-process `dx`
                Dense(3, 32, tanh),                 # Dense Layer 3 -> 32 with `tanh` activation
                Dense(32, 2, tanh),                 # Dense Layer 32 -> 2 with `tanh` activation 
                postProcess,                        # post process `dx`
                dx -> cacheRetrieve(5:6, dx),       # dynamics FMU | dynamics ANN
                gates,                              # compute resulting dx from ANN + FMU
                dx -> cacheRetrieve(1:4, dx))       # stack together: dx[1,2,3,4] from cache + dx[5,6] from gates

    # new NeuralFMU 
    neuralFMU = ME_NeuralFMU(f,                 # the FMU used in the NeuralFMU 
                            model,             # the model we specified above 
                            (tStart, tStop),   # a default start ad stop time for solving the NeuralFMU
                            saveat=tSave)      # the time points to save the solution at
    neuralFMU.modifiedState = false             # speed optimization (NeuralFMU state equals FMU state)

    return neuralFMU 
end

fmu = fmiLoad("VLDM", "Dymola", "2020x"; type=:ME) 
# built the NeuralFMU on basis of the loaded FMU `fmu`
neuralFMU = build_NFMU(fmu)
# a more efficient execution mode
x0 = FMIZoo.getStateVector(data,data.consumption_t[1]) 
fmiSingleInstanceMode(fmu, true)

train_t = data.consumption_t 
train_data = collect([d] for d in data.cumconsumption_val)
BATCHDUR = 4.5
# batch the data (time, targets), train only on model output index 6, plot batch elements
batch = batchDataSolution(neuralFMU,                            # our NeuralFMU model
                        t -> FMIZoo.getStateVector(data, t),  # a function returning a start state for a given time point `t`, to determine start states for batch elements
                        train_t,                              # data time points
                        train_data;                           # data cumulative consumption 
                        batchDuration=BATCHDUR,               # duration of one batch element
                        indicesModel=6:6,                     # model indices to train on (6 equals the state `cumulative consumption`)
                        plot=true,                           # don't show intermediate plots (try this outside of Jupyter)
                        parameters=data.params,               # use the parameters (map file paths) from *FMIZoo.jl*
                        showProgress=showProgress)

ThummeTo / FMIFlux.jl

Plot option in batch.jl function batchDataSolution uses wrong index #117