Scientific machine learning

1 Introduction

In this tutorial, two new approaches to fitting PKPD data will be presented, now that neural networks (NN) have been discussed in the previous course material. The first one is a purely data-driven method, namely a recurrent neural network (RNN), a type of NN architecture focused on sequential data. And the second method embeds a NN into the drug dynamics. This mix of scientific knowledge and AI defines the field of Scientific machine learning (SciML), of which the second approach presented here is an example.

As for the topic of the analysis, we will investigate neutropenia, a medical condition characterized by an abnormally low absolute neutrophil count (ANC). Neutrophils are a type of white cell produced in the bone marrow. They circulate through blood flow and fight bacterial and fungal infections. Neutropenia can be triggered as a side effect of cytotoxic anticancer drugs.

The classical pharmacokinetic-pharmacodynamic (PKPD) model describing the drug-ANC dynamics was proposed in (Friberg et al. 2002). However, researchers still debate the interaction between the proliferating (bone marrow) and circulating (blood) compartments. Therefore, (Soto et al. 2011) compares the original model to 5 alternatives proposed in the literature.

1.1 Environment

As usual, we start by setting up the environment with the required packages and definitions.

using Pumas,
    CairoMakie,
    Statistics,
    Random,
    PumasUtilities,
    DeepPumas,
    Flux,
    Hyperopt,
    Dates,
    ValueHistories
const rng = Pumas.default_rng()
const seed = 500
Random.seed!(rng, seed)

2 Data engineering

2.1 Data generation

A synthetic population will be used. First, the observation times are defined based on the dosage information. More information about representing dosing can be found in the Data Representation in Pumas tutorial. Afterwards, the population is increased by repeating subjects. And doses are determined by the original subject ID multiplied by 1.1e3. This keeps the dataset size divisible by the number of doses, which will be important for data splitting.

n_unique_subjects = 10
# Observations times
obs_times = vcat(range(0, 72, length = 4), range(120, 600, length = 6))
@show permutedims(obs_times)
@show length(obs_times)
repeated_indices = repeat(1:n_unique_subjects; inner = 10)
# Create population
population = map(enumerate(repeated_indices)) do (subject_ID, dose_base)
    Subject(;
        id = string(subject_ID),
        events = DosageRegimen(dose_base * 1.1e3; time = 0.0, cmt = 1, evid = 1),
        time = obs_times,
    )
end;

permutedims(obs_times) = [0.0 24.0 48.0 72.0 120.0 216.0 312.0 408.0 504.0 600.0]
length(obs_times) = 10

Now let’s define the alternative Friberg model. As discussed in the introduction, the drug-ANC interaction is a complex component in the scientific NLME model, being thus a great target for SciML. So, among the alternative Friberg models discussed in (Soto et al. 2011), the one labeled as 4A was chosen here to generate the synthetic data to be used for the remainder of the tutorial. As its defining characteristic, the growth rate of proliferating cells is described by a self-renewal mechanism that is affected by a feedback term, which is inversely proportional to the concentration of circulating neutrophils. According to the authors, the drug inhibited the rate of proliferation as a function of plasma concentration using an inhibitory EMAX model (Soto et al. 2011).

\[ E_{drug} = \frac{Conc}{IC_{50} + Conc} \]

alternative_model = @model begin
    @param begin
        CIRC0 ∈ RealDomain(init = 5.75)
        Rmax ∈ RealDomain(init = 0.0714)
        Rmin ∈ RealDomain(init = 0.0234)
        CL ∈ RealDomain(init = 4)
        VC ∈ RealDomain(init = 70)
        VP ∈ RealDomain(init = 50)
        Q ∈ RealDomain(init = 4)
        KA ∈ RealDomain(init = 1)
        ic50 ∈ RealDomain(init = 13.1)
        km ∈ RealDomain(init = 1.65)
        σ_add ∈ RealDomain(; lower = 1e-6, init = 0.41)
        σ_prop ∈ RealDomain(; lower = 1e-6, init = 0.392)
        σ_add_pd ∈ RealDomain(; lower = 1e-6, init = 0.46)
    end
    @pre begin
        k_tr = Rmax - (Rmax - Rmin) * CIRC0 / (km + CIRC0)
    end
    @init begin
        Prol = CIRC0
        Transit1 = CIRC0
        Transit2 = CIRC0
        Circ = CIRC0
    end
    @vars begin
        conc := Central / VC
        e_drug := conc / (ic50 + conc)
    end
    @dynamics begin
        # @Central1Periph1
        Depot' = -KA * Depot
        Central' = KA * Depot - (CL + Q) / VC * Central + Q / VP * Peripheral
        Peripheral' = Q / VC * Central - Q / VP * Peripheral

        Prol' =
            (Rmax - (Rmax - Rmin) * Circ / (km + Circ)) * (1 - e_drug) * Prol - k_tr * Prol
        Transit1' = k_tr * Prol - k_tr * Transit1
        Transit2' = k_tr * Transit1 - k_tr * Transit2
        Circ' = k_tr * Transit2 - k_tr * Circ
    end
    @derived begin
        cp := @. Central / VC
        pk ~ @. Normal(cp, sqrt(σ_add^2 + (cp * σ_prop)^2))
        pd ~ @. Normal(Circ, σ_add_pd)
    end
end

PumasModel
  Parameters: CIRC0, Rmax, Rmin, CL, VC, VP, Q, KA, ic50, km, σ_add, σ_prop, σ_add_pd
  Random effects:
  Covariates:
  Dynamical system variables: Depot, Central, Peripheral, Prol, Transit1, Transit2, Circ
  Dynamical system type: Nonlinear ODE
  Derived: pk, pd
  Observed: pk, pd

The alternative Friberg model is used to simulate PK and PD observations in the population just created.

const references = init_params(alternative_model)
# Simulate from alternative Friberg model
simulated_population =
    Subject.(simobs(alternative_model, population, references; rng, obstimes = obs_times))

Population
  Subjects: 100
  Observations: pk, pd

2.2 NLME diagnostic tools

To examine the model, Pumas provides many diagnostic utilities. For certain plots, fitting the model is required, so the next cell fits alternative_model to the data simulated from it. More information on diagnostic plots can be found in the Model Diagnostics and Evaluation in Pumas tutorial.

Note that these diagnostics are all “in-sample”; they say nothing about how well a Pumas model performs when compared with new clinical data not seen in training.

fpm_alternative = fit(
    alternative_model,
    simulated_population,
    references,
    NaivePooled();
    optim_options = (; show_trace = false),
    checkidentification = false,
)

One option for diagnostics is to simulate from the fitted model at a finer time resolution with simobs and plot the predicted circulation ANC over time for each subject, color-coded by dose.

# Simulate with higher time resolution
fine_time = 0.0:0.1:obs_times[end]
plot_sims = simobs(
    alternative_model,
    Subject.(simulated_population),
    coef(fpm_alternative);
    obstimes = fine_time,
)
# Figure and axis
figure = Figure(; size = (800, 800), fontsize = 20)
axis = Axis(figure[1, 1]; xlabel = "Time", ylabel = "Circulation ANC (PD)")
# Pre-compute one distinct color per subject
subject_colors = cgrad(:oxy, length(simulated_population); categorical = true)
# Per subject, extract information and plot PD trend
(; Rmax, Rmin, km) = coef(fpm_alternative)
for (i, target_sub) in enumerate(plot_sims)
    c = subject_colors[i]
    circ = target_sub.dynamics.Circ
    label = target_sub.subject.events[1].amt |> Int |> string
    trend = @. (Rmax - (Rmax - Rmin) * circ / (km + circ))
    lines!(axis, fine_time, trend; color = c, label)
end
axislegend(axis, "Doses"; position = :rt, unique = true)
display(figure);

inspect() is called on the FittedPumasModel to create a summary of the model predictions, residuals, Empirical Bayes estimates (not applied here), etc.

insp = inspect(fpm_alternative);

[ Info: Calculating predictions.
[ Info: Calculating weighted residuals.
[ Info: Calculating empirical bayes.
[ Info: Evaluating dose control parameters.
[ Info: Evaluating individual parameters.
[ Info: Done.

A ready-made plot is the convergence_trace(), including the objective and the gradient’s norm over iterations.

figure = (; size = (800, 800), fontsize = 18)
display(convergence_trace(fpm_alternative; figure));

And the classical (goodness-of-fit) is also provided. It plots observed values against both population predictions (PRED, from the fixed effects alone) and individual predictions (IPRED, conditioned on empirical Bayes estimates of the random effects). Points clustering tightly around the identity line indicate that the model’s predicted means track the data well; systematic curvature or spread reveals bias or misspecification in the structural model. Note that population and individual plots are the same here, since there are no random effects.

display(goodness_of_fit(insp; figure));

The distributions of weighted residuals (docs) can be seen in the following plot. Weighted residuals (WRES) are raw residuals rescaled by the model-predicted standard deviation, so they should be approximately standard-normal if the error model is correct. The plot shows their distribution against PREDs and IPREDs (with no difference between the two here, because there are no random effects in the model).

display(wresiduals_dist(insp; figure));

Finally, a visual predictive check (VPC) of PD is built in the cell block below. The VPC is a simulation-based diagnostic: vpc() repeatedly simulates from the fitted model and, at each observation time, collects the distribution of the simulated outputs. vpc_plot() then overlays the resulting prediction intervals (typically the 10th–90th percentile band and the median) on the observed data percentiles. When the model is well specified, the observed quantiles should fall within the corresponding simulated bands.

display(vpc_plot(alternative_model, vpc(fpm_alternative; observations = [:pd]); figure));

[ Info: Continuous VPC

2.3 Data pre-processing

Turning now to the data layout, both the inputs and outputs defined below contain three dimensions: variables (rows), time points (columns), subjects (depth). For inputs, the variables are time deltas and dose; for outputs, PK and PD trends.

trend_size = length(obs_times)
population_size = length(simulated_population)
# Reshape data into three-dimensional arrays
inputs = zeros(Float32, (2, trend_size, population_size))
outputs = zeros(Float32, (2, trend_size, population_size))
for (subject_ID, subject) in enumerate(simulated_population)
    # Time deltas
    inputs[1, :, subject_ID] .= [0; diff(subject.time)]
    # Time profile of dose
    inputs[2, :, subject_ID] .= vcat([subject.events[1].amt], zeros(trend_size - 1))
    # PK trend
    outputs[1, :, subject_ID] .= subject.observations.pk
    # PD trend
    outputs[2, :, subject_ID] .= subject.observations.pd
end

The dataset is split according to doses as the following table shows. Each split is itself a three-dimensional array with the layout described previously.

Split	Dose	Number of subjects
training	{1 \(\dots\) 8} \(\cdot\) 1100	80
validation	9900	10
test	11000	10

doses = sort(unique(inputs[2, 1, :]))
splits = Dict{Symbol,AbstractArray{Float32,3}}()
# Training
training_doses = doses[1:8]
training_indices = findall(in(training_doses), @view inputs[2, 1, :])
splits[:training_inputs] = @view inputs[:, :, training_indices]
splits[:training_outputs] = @view outputs[:, :, training_indices]
# Validation
validation_indices = findall(==(doses[9]), @view inputs[2, 1, :])
splits[:validation_inputs] = @view inputs[:, :, validation_indices]
splits[:validation_outputs] = @view outputs[:, :, validation_indices]
# Test
test_indices = findall(==(doses[10]), @view inputs[2, 1, :])
splits[:test_inputs] = @view inputs[:, :, test_indices]
splits[:test_outputs] = @view outputs[:, :, test_indices]

Inputs are normalized by (x .- training_minimum) / training_std per variable. Outputs are kept in their original scale; the maximum training target values are used to rescale RNN predictions in the loss function.

# Normalization parameters for inputs
minimum_time_delta = minimum(@view splits[:training_inputs][1, :, :])
deviation_time_delta = std(@view splits[:training_inputs][1, :, :])
minimum_dose = minimum(@view splits[:training_inputs][2, :, :])
deviation_dose = std(@view splits[:training_inputs][2, :, :])
# Maximum target values for rescaling RNN predictions
max_PK = Float32(maximum(@view splits[:training_outputs][1, :, :]))
max_PD = Float32(maximum(@view splits[:training_outputs][2, :, :]))
const max_targets = [max_PK; max_PD]
# Transformation of inputs
for split_input in [:training_inputs, :validation_inputs, :test_inputs]
    @. splits[split_input][1, :, :] =
        (splits[split_input][1, :, :] - minimum_time_delta) / deviation_time_delta
    @. splits[split_input][2, :, :] =
        (splits[split_input][2, :, :] - minimum_dose) / deviation_dose
end
# Check results
for (split_name, data) in splits, row = 1:2
    q = quantile(data[row, :, :], [0.0, 0.25, 0.5, 0.75, 1.0])
    println("$split_name row $row  ", round.(q; digits = 2))
end

test_inputs row 1  [0.0, 0.64, 1.92, 2.55, 2.55]
test_inputs row 2  [0.0, 0.0, 0.0, 0.0, 6.52]
training_outputs row 1  [-1.07, -0.05, 0.39, 2.39, 41.8]
training_outputs row 2  [0.74, 4.26, 5.4, 6.12, 8.57]
validation_inputs row 1  [0.0, 0.64, 1.92, 2.55, 2.55]
validation_inputs row 2  [0.0, 0.0, 0.0, 0.0, 5.87]
test_outputs row 1  [-1.01, -0.03, 0.57, 8.14, 45.6]
test_outputs row 2  [0.95, 3.29, 5.33, 6.44, 8.26]
training_inputs row 1  [0.0, 0.64, 1.92, 2.55, 2.55]
training_inputs row 2  [0.0, 0.0, 0.0, 0.0, 5.22]
validation_outputs row 1  [-1.25, -0.09, 0.4, 7.06, 54.4]
validation_outputs row 2  [0.81, 3.48, 5.57, 6.42, 8.44]

3 RNN

RNNs are a type of neural network architecture designed to handle time-series data, such as predicting the weather in a sequence of days. In this tutorial, we apply RNNs and dense neural layers to predict PK and PD trends from two inputs: time intervals between observations, and dose at each time point.

With the data ready, a few definitions are required to setup training and keep the code organized. The function setup_training builds the architecture of the NN model, including an RNN layer and a varying number of dense layers. The depth (number of layers - network_depth) and width (size of each layer - layer_multiplier) of the model are among the hyperparameters to be optimized. In the end, the trainable values are returned, later called learned.

output_dimension = size(splits[:training_outputs], 1)
function setup_training(hyperparameters)
    (; layer_multiplier, network_depth, activation) = hyperparameters
    # Reset seed so, during hyperparameter optimization, all iterations
    # start with the same model
    Random.seed!(rng, seed)
    # Variable model architecture
    M = output_dimension * layer_multiplier
    hidden_layers = [Dense(M => M, activation) for _ = 1:network_depth]
    rnn = Chain(
        # RNN() has default activation of tanh
        RNN(size(splits[:training_inputs], 1) => M, activation),
        hidden_layers...,
        Dense(M => output_dimension),
    )
    # Sample large initial values for residual error parameters
    σ_add = references.σ_add * 3
    σ_prop = references.σ_prop * 3
    σ_add_pd = references.σ_add_pd * 3
    # NamedTuple concentrating trainable elements
    return (;
        rnn_model = rnn,
        log_σ_add = [log(σ_add)],
        log_σ_prop = [log(σ_prop)],
        log_σ_add_pd = [log(σ_add_pd)],
    )
end

At each iteration of the hyperparameter optimization, the proposed combination of hyperparameters is experimented with by means of the early-stopping setup defined in early_stopping!. Another two hyperparameters are used here when creating OptimiserChain():

1. Gradient clipping: limit size of gradient to avoid the problem of exploding gradients
2. Learning rate: try different step sizes to train the NN model

es = Flux.early_stopping(f, delay; distance = -, init_score, min_dist) (docs) sets the details of early-stopping with the following definitions:

1. f: function that returns scalar loss (here, on the validation split)
2. delay: how many consecutive times f can return a value that’s not better than the best so far. After this, training is stopped by es() && break
3. distance: indicate if the value returned by f should be minimized (-) or maximized (+)
4. init_score: value to be considered as a reference in the beginning. It will be used for comparison when es() is called for the first time
5. min_dist: minimal distance between the value returned by f and the current best score to consider it an improvement

Note that es is itself a function, with the same arguments as f. In our case, there will be zero arguments.

A strict limit of max_epochs epochs is specified for training to avoid cases that improve almost asymptotically, taking a long time to trigger early-stopping while resulting in a bad fit. Still inside the while loop, other than actual training, test and validation performance is also logged to history for later plotting.

const validation_inputs = splits[:validation_inputs]
const validation_outputs = splits[:validation_outputs]
const test_inputs = splits[:test_inputs]
const test_outputs = splits[:test_outputs]
function early_stopping!(
    data_loader,
    learned,
    history,
    hyperparameters,
    loss;
    max_epochs = 100,
)
    (; gradient_clipping, learning_rate, layer_multiplier, network_depth, activation) =
        hyperparameters
    # Setup optimizer
    opt = OptimiserChain(ClipNorm(gradient_clipping), Optimisers.Adam(learning_rate))
    optimizer_state = Flux.setup(opt, learned)
    # Early-stopping (reads from val_loss, refreshed each epoch below)
    val_loss = Ref{Float32}(Inf)
    es = Flux.early_stopping(() -> val_loss[], 160; min_dist = 0, init_score = Inf)
    epoch = 1
    # Main loop
    while epoch <= max_epochs
        # Train in batches
        batch_loss = map(data_loader) do (batch_input, batch_output)
            loss_value, grad =
                Flux.withgradient(r -> loss(r, batch_input, batch_output), learned)
            Flux.update!(optimizer_state, learned, grad[1])
            return loss_value
        end
        val_loss[] = loss(learned, validation_inputs, validation_outputs)
        # Log validation and test performances
        if epoch < 5 || epoch % 10 == 0
            push!(history, :train, epoch, sum(batch_loss))
            push!(history, :val, epoch, val_loss[])
            push!(history, :test, epoch, loss(learned, test_inputs, test_outputs))
        end
        es() && break
        epoch += 1
    end
    return epoch
end

Note

In Julia, a Ref object is a container with exactly one element, which is set, modified and accessed as shown in the following example:

tinybox = Ref{Int}(42)

Base.RefValue{Int64}(42)

tinybox[]

42

tinybox[] = 7

7

tinybox[]

7

A common use case for Ref is in multi-argument broadcasting, where you want to regard the wrapped object as a “scalar” (a one-element container). We have:

vcat.(Ref(1:2), [4:5, 5:6])

2-element Vector{Vector{Int64}}:
 [1, 2, 4, 5]
 [1, 2, 5, 6]

but

vcat.(1:2, [4:5, 5:6]) # 1:2 regarded as a two-element container

2-element Vector{Vector{Int64}}:
 [1, 4, 5]
 [2, 5, 6]

Later we will encounter the following use case:

t = (; x = 1, y = 2)
getfield.(Ref(t), (:y, :x))

(2, 1)

Remove the Ref(…) wrapper and an exception is thrown.

Let’s define function to calculate error in NN model prediction. The RNN predictions are used as means to build normal distributions around each PK and PD time point, with the learned sigmas defining the standard deviations. Then the negative log of the probabilities of labels given these distributions are averaged.

function loss(learned, input, output)::Float32
    (; rnn_model, log_σ_add, log_σ_prop, log_σ_add_pd) = learned
    σ_add, σ_prop, σ_add_pd = exp.(first.((log_σ_add, log_σ_prop, log_σ_add_pd)))
    # Rescale to original values
    rnn_output = rnn_model(input) .* max_targets
    # Cumulative log-loss across subjects
    cumulative_log_loss = sum(axes(rnn_output, 3)) do subject_ID
        # Extract both trends for current subject
        pk_preds = @view rnn_output[1, :, subject_ID]
        pd_preds = @view rnn_output[2, :, subject_ID]
        # Negative log-probabilities of labels given distributions learned by RNN
        ll = zero(eltype(pk_preds))
        for j in axes(pk_preds, 1)
            # Build normal distributions around each PK and PD observation
            ll -= logpdf(
                Normal(pk_preds[j], sqrt(σ_add^2 + (pk_preds[j] * σ_prop)^2)),
                output[1, j, subject_ID],
            )
            ll -= logpdf(Normal(pd_preds[j], σ_add_pd), output[2, j, subject_ID])
        end
        return ll
    end
    return cumulative_log_loss
end

Note

Internally, Flux.jl uses Float32 precision. Hence the ::Float32 in the function definition specifying the type of the output. Without this, the output would be in 64-bit precision, Flux.jl would trigger a warning.

The following code block specifies some initializations for the hyperparameter optimization. First, the dataloader, which prepares the data to be iterated upon in batches. Then, log_DF will store, for each iteration: system time, the 5 hyperparameters tried, final validation performance during early-stopping, and total epochs used in that run. Next, inner is a vector articulating, for hyperband optimization (see next section), the type of sampling to be applied to each hyperparameter (all categorical). Of these, activations switches the activation function used in the hidden layers of the NN model among ReLU, sigmoid (σ), and hyperbolic tangent (tanh).

const data_loader =
    Flux.DataLoader((splits[:training_inputs], splits[:training_outputs]); batchsize = 20)
log_DF = DataFrame([
    :Time => "0",
    :gradient_clipping => 0.0,
    :learning_rate => 0.0,
    :Layer_multiplier => 0,
    :network_depth => 0,
    :activation => relu,
    :Final_validation => 0.0,
    :Total_epochs => 1,
])
pop!(log_DF)
inner = Hyperopt.BOHB(
    dims = [
        Hyperopt.Categorical(4),
        Hyperopt.Categorical(3),
        Hyperopt.Categorical(4),
        Hyperopt.Categorical(4),
        Hyperopt.Categorical(3),
    ],
)
activations = [relu, σ, tanh]

main defines the core steps that will take place every iteration of the hyperparameter optimization. After setup_training(), history = MVHistory() is created to log out-of-sample performance. After the main training steps with early_stopping!(), some information is returned: how many epochs training took, the performance trends in history and the learned named tuple with the NN model. Also, livePlot() (see appendix Section 7.1) can be optionally called to plot the current model’s performance in all splits, showing the early-stopping history of the run. total_budget will be explained in the next section.

total_budget = 350
function main(hyperparameters; max_epochs)
    learned = setup_training(hyperparameters)
    history = MVHistory()
    # Train with early-stopping
    epoch =
        early_stopping!(data_loader, learned, history, hyperparameters, loss; max_epochs)
    # Uncomment next line for optional plotting; see appendix:
    # epoch >= total_budget * 0.4 && livePlot(history, hyperparameters)
    return epoch, history, learned
end

3.1 Hyperparameter optimization

In previous tutorials, hyperparameter optimization mainly used random search. That’s mostly a default to try different options. Without diving into the field of hyperparameter optimization (Bartz et al. 2023), we’ll use a different algorithm provided by the package Hyperopt.jl, Hyperband (Li et al. 2018).

It tries to distribute a finite resource among experiments with different hyperparameter configurations. In the for loop macro below, resources is the total budget, which will be spread among multiple attempts. The definition of resource is context-dependent and chosen by the researcher. Here it’s the limit of epochs for the early-stopping loop. R is the maximum amount of resource that can be allocated to a single configuration. And \(\eta\) controls the rough proportion of configurations discarded in each round of “successive halving”, when the algorithm discards some less promising experiments and redistributes resources to runs with better performances.

hyperopt_result = @hyperopt for resources in total_budget,
    sampler in Hyperband(R = total_budget * 0.4, η = 3; inner),
    gradient_clipping in exp10.(-8:3:1),
    learning_rate in exp10.(-3:-1),
    layer_multiplier = 1:5:16,
    network_depth = 1:4:13,
    activation = 1:3

    # Package API
    if !(state === nothing)
        gradient_clipping, learning_rate, layer_multiplier, network_depth, activation =
            state
    end
    hyperparameters = (;
        gradient_clipping,
        learning_rate,
        layer_multiplier,
        network_depth,
        activation = activations[activation],
    )
    # Execute core training procedures
    epoch_limit, history, learned = main(hyperparameters; max_epochs = resources)
    # Log current attempt to log_DF
    final_validation = round(get(history, :val)[2][end]; digits = 6)
    push!(
        log_DF,
        [
            string(now())[12:(end-4)],
            gradient_clipping,
            learning_rate,
            layer_multiplier,
            network_depth,
            activations[activation],
            final_validation,
            epoch_limit - 1,
        ];
        promote = true,
    )
    # Hyperopt.jl API: return metric and associated hyperparameters
    return final_validation,
    (
        gradient_clipping,
        learning_rate,
        layer_multiplier,
        network_depth,
        activation,
        learned,
    )
end;

After optimization, print logs with the 10 best runs and compare sigmas.

sort!(log_DF, :Final_validation)
println("Top 10 configurations:")
show(first(log_DF, min(10, nrow(log_DF))); allcols = true)
# Extract fitted residual error parameters
best_learned = hyperopt_result.minimizer[end]
σ_add_fit = exp(only(best_learned.log_σ_add))
σ_prop_fit = exp(only(best_learned.log_σ_prop))
σ_add_pd_fit = exp(only(best_learned.log_σ_add_pd))
true_sigmas = references[[:σ_add, :σ_prop, :σ_add_pd]] |> values |> collect
sigma_zip = zip([σ_add_fit, σ_prop_fit, σ_add_pd_fit], true_sigmas)
deviations = map(sigma_zip) do (fitted, target)
    return round((fitted / target - 1) * 100; digits = 1)
end
println(
    "\n\nResidual error:\n",
    DataFrame(
        Parameter = [:σ_add, :σ_prop, :σ_add_pd],
        Fitted_values = [σ_add_fit, σ_prop_fit, σ_add_pd_fit],
        True_values = true_sigmas,
        Dev_percent = deviations,
    ),
)

Top 10 configurations:
10×8 DataFrame
 Row │ Time      gradient_clipping  learning_rate  Layer_multiplier  network_depth  activation  Final_validation  Total_epochs
     │ String    Float64            Float64        Int64             Int64          Function    Float64           Int64
─────┼─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────
   1 │ 20:29:40            10.0              0.01                11              5  relu                 209.145           140
   2 │ 20:28:49            10.0              0.01                11              5  relu                 225.387            46
   3 │ 20:32:04             0.01             0.01                11              9  relu                 230.236           140
   4 │ 20:36:46             2.0              0.01                 3              3  relu                 247.145           140
   5 │ 20:37:51             1.0              0.01                 4              4  relu                 274.476           140
   6 │ 20:39:25             1.0              0.01                 4              4  relu                 274.476           140
   7 │ 20:37:18             1.0e-5           0.01                 6              5  tanh                 300.025           140
   8 │ 20:33:29            10.0              0.01                 6              5  tanh                 306.723            46
   9 │ 20:34:18            10.0              0.01                 6              5  tanh                 312.49            140
  10 │ 20:33:39             4.0              0.01                 4              4  tanh                 328.325            46

Residual error:
3×4 DataFrame
 Row │ Parameter  Fitted_values  True_values  Dev_percent
     │ Symbol     Float64        Float64      Float64
─────┼────────────────────────────────────────────────────
   1 │ σ_add           0.450442        0.41           9.9
   2 │ σ_prop          0.398267        0.392          1.6
   3 │ σ_add_pd        0.458931        0.46          -0.2

3.2 Diagnosis

To judge the out-of-sample performance of the best RNN, an observations-predictions plot is drawn for PK and PD.

# Choose split and extract data
diagnose_split = :test
diagnose_inputs = splits[Symbol(diagnose_split, "_inputs")]
diagnose_outputs = splits[Symbol(diagnose_split, "_outputs")]
# Get predictions and best model
rnn_model = hyperopt_result.minimizer[end][1]
predictions = rnn_model(diagnose_inputs) .* max_targets
# Plot
fig = Figure(; size = (800, 800), fontsize = 22)
for (col, title) in enumerate(("PK", "PD"))
    # Prepare data for plotting
    targets = @view diagnose_outputs[col, :, :]
    preds = @view predictions[col, :, :]
    plot_info = zip(targets, preds) |> collect |> vec
    start_line = min(minimum(preds), minimum(targets))
    finish_line = max(maximum(preds), maximum(targets))
    # Create Axis and plot
    ylabel = col == 1 ? "Predictions" : ""
    axis = Axis(fig[1, col]; title, ylabel, xlabel = "Observations")
    scatter!(axis, plot_info; marker = 'o')
    lines!(
        axis,
        [(start_line, start_line), (finish_line, finish_line)];
        color = :black,
        linewidth = 2,
    )
end
display(fig);

Observations and predictions can also be plotted over time, as done in the next code block. The plots are stratified by dose, including validation (9.9e3) and test (1.1e4) splits. The fitted residual-error parameters are used to draw confidence bands as well.

n_doses = population_size ÷ n_unique_subjects
# Setup figure and labels
fig = Figure(; size = (900, 170 * n_doses), fontsize = 22)
Label(fig[0, 1], "PK"; font = :bold, tellwidth = false)
Label(fig[0, 2], "PD"; font = :bold, tellwidth = false)
row = 0
for diagnose_split in [:training, :validation, :test]
    # For split, get and use data
    diagnose_inputs = splits[Symbol(diagnose_split, "_inputs")]
    diagnose_outputs = splits[Symbol(diagnose_split, "_outputs")]
    predictions = rnn_model(diagnose_inputs) .* max_targets
    # Doses in current split
    diagnose_doses = unique(@view diagnose_inputs[2, 1, :])
    # Per dose, plot observations and predictions for PK and PD
    for dose in diagnose_doses
        row += 1
        # Get subjects and dose value before normalization
        subjects_current_dose = findall(==(dose), @view diagnose_inputs[2, 1, :])
        original_dose = (dose * deviation_dose) + minimum_dose
        dose_label = "Dose " * string(round(Int, original_dose))
        # Create axes for current row
        PK_axis = Axis(fig[row, 1]; ylabel = dose_label)
        ylims!(PK_axis, (0.0, max_targets[1]))
        PD_axis = Axis(fig[row, 2])
        ylims!(PD_axis, (0.0, max_targets[2]))
        # Suppress x axis details in all but the last row
        if row != n_doses
            hidexdecorations!(PK_axis)
            hidexdecorations!(PD_axis)
        else
            PK_axis.xlabel = PD_axis.xlabel = "Time"
        end
        # Iterate in subjects receiving current dose
        for subject_ID in subjects_current_dose
            PK_axis.xgridvisible = PD_axis.xgridvisible = true
            # PK observations
            scatter!(
                PK_axis,
                obs_times,
                @view diagnose_outputs[1, :, subject_ID];
                alpha = 0.5,
                color = :darkgreen,
            )
            # PK RNN predictions
            pk_pred = predictions[1, :, subject_ID]
            lines!(PK_axis, obs_times, pk_pred; linewidth = 2, color = :black)
            # PK error band (±(2*SD) from combined additive-proportional error model)
            pk_std = @. sqrt(σ_add_fit^2 + (pk_pred * σ_prop_fit)^2)
            band!(
                PK_axis,
                obs_times,
                pk_pred .- 2 .* pk_std,
                pk_pred .+ 2 .* pk_std;
                color = (:purple, 0.03),
            )
            # PD observations
            scatter!(
                PD_axis,
                obs_times,
                @view diagnose_outputs[2, :, subject_ID];
                alpha = 0.5,
                color = :darkgreen,
                label = "Observations",
            )
            # PD RNN predictions
            pd_pred = predictions[2, :, subject_ID]
            lines!(
                PD_axis,
                obs_times,
                pd_pred;
                linewidth = 2,
                color = :black,
                label = "Predictions",
            )
            # PD error band (±(2*SD) from additive error model)
            band!(
                PD_axis,
                obs_times,
                pd_pred .- 2 .* σ_add_pd_fit,
                pd_pred .+ 2 .* σ_add_pd_fit;
                color = (:purple, 0.03),
            )
        end
        if row == 1
            axislegend(PD_axis; position = :rb, unique = true) # Include legend once
        elseif row == 9
            Label(
                fig[row, 0],
                "Validation";
                font = :bold,
                tellheight = false,
                rotation = π / 2,
            )
        elseif row == 10
            Label(fig[row, 0], "Test"; font = :bold, tellheight = false, rotation = π / 2)
        end
    end
end
display(fig);

4 SciML

The inclusion of AI in scientific analysis isn’t “all or nothing”. SciML is a great example, focusing on models that mix mechanistic knowledge about laws of nature with data-driven ML models. So far, an alternative to the mechanistic Friberg model was used to generate a synthetic dataset. Then, a purely data driven RNN model was used to fit the simulated observations.

In this section, the same will be done using a SciML universal differential equation (UDE) model, which embeds a NN in an NLME model. The fundamental difference is in the rate of change of the proliferation compartment, where the NN (a multilayer perceptron called mlp below) tunes the rate of generation of progenitor cells in Prol.

\[ Prol' = k_{tr} \cdot \textbf{mlp}\left(\frac{Circ}{scale_{circ}}\right) \cdot (1 - e_{drug}) \cdot Prol - k_{tr} \cdot Prol \]

For comparison, the following was the formulation used in the alternative Friberg model, generator of the data.

\[ Prol' = \left(R_{max} - (R_{max} - R_{min}) * \frac{Circ}{k_m + Circ}\right)* (1 - e_{drug}) * Prol - k_{tr} * Prol \]

ude_model = @model begin
    @param begin
        mlp ∈ MLPDomain(1, 3, 3, (1, softplus))
        scale_circ ∈ RealDomain(init = 1)
        CIRC0 ∈ RealDomain(; lower = 1e-3, init = 5.75)
        Rmax ∈ RealDomain(; lower = 1e-3, init = 0.0714)
        Rmin ∈ RealDomain(; lower = 1e-3, init = 0.0234)
        CL ∈ RealDomain(; lower = 1e-3, init = 4)
        VC ∈ RealDomain(; lower = 1e-3, init = 70)
        VP ∈ RealDomain(; lower = 1e-3, init = 50)
        Q ∈ RealDomain(; lower = 1e-3, init = 4)
        KA ∈ RealDomain(; lower = 1e-3, init = 1)
        ic50 ∈ RealDomain(; lower = 1e-3, init = 13.1)
        km ∈ RealDomain(; lower = 1e-3, init = 1.65)
        σ_add ∈ RealDomain(; lower = 1e-6, init = 0.41)
        σ_prop ∈ RealDomain(; lower = 1e-6, init = 0.392)
        σ_add_pd ∈ RealDomain(; lower = 1e-6, init = 0.46)
    end
    @pre begin
        k_tr = Rmax - (Rmax - Rmin) * CIRC0 / (km + CIRC0)
    end
    @init begin
        Prol = CIRC0
        Transit1 = CIRC0
        Transit2 = CIRC0
        Circ = CIRC0
    end
    @vars begin
        conc := Central / VC
        e_drug := conc / (ic50 + conc)
    end
    @dynamics begin
        # @Central1Periph1
        Depot' = -KA * Depot
        Central' = KA * Depot - (CL + Q) / VC * Central + Q / VP * Peripheral
        Peripheral' = Q / VC * Central - Q / VP * Peripheral

        Prol' = k_tr * mlp(Circ / scale_circ)[1] * (1 - e_drug) * Prol - k_tr * Prol
        Transit1' = k_tr * Prol - k_tr * Transit1
        Transit2' = k_tr * Transit1 - k_tr * Transit2
        Circ' = k_tr * Transit2 - k_tr * Circ
    end
    @derived begin
        cp := @. Central / VC
        pk ~ @. Normal(cp, sqrt(σ_add^2 + (cp * σ_prop)^2))
        pd ~ @. Normal(Circ, σ_add_pd)
    end
end

PumasModel
  Parameters: mlp, scale_circ, CIRC0, Rmax, Rmin, CL, VC, VP, Q, KA, ic50, km, σ_add, σ_prop, σ_add_pd
  Random effects:
  Covariates:
  Dynamical system variables: Depot, Central, Peripheral, Prol, Transit1, Transit2, Circ
  Dynamical system type: Nonlinear ODE
  Derived: pk, pd
  Observed: pk, pd

It’s important to highlight that the scope of the NN inside the model can be controlled. For example, it could represent the entire generation of progenitor cells, instead of just tuning the rate. There could even be multiple NNs in the model. This type of control can be used to limit overfitting, since it balances the amount of knowledge and data-driven components in the model.

Fitting of the model will include multistart, multiple fits starting from random initializations of the parameters values. sample_ude_parameters() will sample the initial parameter values, where param_range indicates the range around the true values: \(true\_value \cdot (1 \pm param\_range)\). This range applies to all parameters except the ones in mlp.

const param_range = 0.4
function sample_ude_parameters()
    mlp = sample_params(ude_model).mlp
    inits = init_params(ude_model)
    non_nn = map(inits[Not(:mlp)]) do value
        lo = (1 - param_range) * value
        hi = (1 + param_range) * value
        return lo + (hi - lo) * rand(rng)
    end
    return (; mlp, non_nn...)
end

Each time we fit the UDE model, we will use the same small subset of the training subjects. Optimization consists of selecting the fitted model having the lowest negative log-likelihood on the test subjects.

# Sample population
ude_subjects = simulated_population[1:10:80]
starting_parameters = []
# Fit UDE model with multistart
multistarts = 3
iterations = 60
fits = map(1:multistarts) do i
    println("\nMultistart $i")
    # Sample and store initial parameter values
    start = sample_ude_parameters()
    push!(starting_parameters, start)
    fpm_ude = fit(
        ude_model,
        ude_subjects,
        start,
        NaivePooled();
        optim_options = (; iterations, show_every = Int(iterations / 2)),
        diffeq_options = (;
            abstol = 1e-14,
            reltol = 1e-14,
            alg = Rodas5P(),
            maxiters = 16e5,
        ),
        checkidentification = false,
    )
    # compute the negative log-likelihood of the validation subjects:
    log_loss =
        -loglikelihood(
            ude_model,
            simulated_population[validation_indices],
            coef(fpm_ude),
            NaivePooled(),
        )
    return (; fpm_ude, log_loss)
end
# Find best fit based on simulation RMSE
fpm_ude = argmin(f -> f.log_loss, fits)[:fpm_ude];

Multistart 1
Iter     Function value   Gradient norm 
     0     2.583319e+03     4.674320e+03
 * time: 0.22991514205932617
    30     2.088188e+02     6.103883e+01
 * time: 338.76199197769165
    60     1.316736e+02     3.685724e+01
 * time: 716.4322929382324

Multistart 2
Iter     Function value   Gradient norm 
     0     4.647630e+03     9.043327e+03
 * time: 0.00017905235290527344
    30     2.257789e+02     2.639217e+01
 * time: 670.7110810279846
    60     2.175539e+02     1.511761e+02
 * time: 993.7718179225922

Multistart 3
Iter     Function value   Gradient norm 
     0     3.438637e+03     6.408193e+03
 * time: 7.605552673339844e-5
    30     2.333823e+02     8.170565e+01
 * time: 391.97560596466064
    60     1.441216e+02     1.416565e+02
 * time: 852.1769881248474

As we did for the RNN model, let’s compare the learned error parameters with their true values:

optimal_parameters = coef(fpm_ude)
σ_add_fit, σ_prop_fit, σ_add_pd_fit =
    getfield.(Ref(optimal_parameters), (:σ_add, :σ_prop, :σ_add_pd))
sigma_zip = zip([σ_add_fit, σ_prop_fit, σ_add_pd_fit], true_sigmas)
deviations = map(sigma_zip) do (fitted, target)
    return round((fitted / target - 1) * 100; digits = 1)
end
println(
    "\n\nResidual error:\n",
    DataFrame(
        Parameter = [:σ_add, :σ_prop, :σ_add_pd],
        Fitted_values = [σ_add_fit, σ_prop_fit, σ_add_pd_fit],
        True_values = true_sigmas,
        Dev_percent = deviations,
    ),
)

Residual error:
3×4 DataFrame
 Row │ Parameter  Fitted_values  True_values  Dev_percent
     │ Symbol     Float64        Float64      Float64
─────┼────────────────────────────────────────────────────
   1 │ σ_add           0.390668        0.41          -4.7
   2 │ σ_prop          0.440264        0.392         12.3
   3 │ σ_add_pd        0.409401        0.46         -11.0

4.1 Diagnosis

As for the true model, multiple plots allow us to take a closer look at the fit, namely the best among all multistarts. Again, this starts with a simulation from the fitted model using a higher time resolution.

fitted_parameters = coef(fpm_ude)
ude_eval_subjects = simulated_population[[80, 90, 100]]
UDE_sims = simobs(ude_model, ude_eval_subjects, fitted_parameters; obstimes = fine_time)

Simulated population (Vector{<:Subject})
  Simulated subjects: 3
  Simulated variables: pk, pd

Let’s track the compartments and NN over time. The next code block builds a figure with an axis for concentrations in circulation (circ), proliferation (prol) and central (Central) compartments. Additionally, NN outputs are plotted to make sure the tanh activations in the hidden layers aren’t saturated, and that the NN output has reasonable magnitude.

# Figure and axes
fig = Figure(; size = (800, 1200), fontsize = 20)
limits = ((nothing, nothing), (0.0, nothing))
circ_axis = Axis(fig[1, 1]; title = "Circ", xlabel = "Time", limits)
prol_axis = Axis(fig[1, 2]; title = "Prol", xlabel = "Time", limits)
nn_axis = Axis(fig[2, 2]; title = "NN", xlabel = "circ / scale_circ")
conc_axis = Axis(fig[2, 1]; title = "Conc", xlabel = "Time")
nn_friberg_axis = Axis(
    fig[3, 2];
    xlabel = "Friberg term replaced by NN",
    ylabel = "NN",
    title = "Comparing NN with Alt. Friberg",
)
# Access fitted parameters in SciML model
(; mlp, Rmax, Rmin, CIRC0, km, VC, scale_circ) = fitted_parameters
k_tr = Rmax - (Rmax - Rmin) * CIRC0 / (km + CIRC0)
nn_ins = Array{Float32}[]
for (sim_sub, target_sub) in zip(UDE_sims, ude_eval_subjects)
    circ = sim_sub.dynamics.Circ # Extract values from subject
    # Plot trends in compartments
    lines!(circ_axis, fine_time, circ)
    lines!(prol_axis, fine_time, sim_sub.dynamics.Prol)
    lines!(
        conc_axis,
        fine_time,
        sim_sub.dynamics.Central / VC; # conc = Central / VC
        label = sim_sub.subject.events[1].amt |> Int |> string,
    )
    # Scale inputs by fitted scale_circ and use fitted NN
    nn_input = Float32.(circ ./ scale_circ)
    push!(nn_ins, nn_input)
    nn_output = only.(mlp(nn_input'))  # Pass as 1×N matrix. Extract scalars
    lines!(
        nn_axis,
        nn_input,
        vec(nn_output) .* k_tr;
        alpha = 0.2,
        linewidth = 3,
        color = :black,
    )
    # get term from Alt. Friberg that NN replaces:
    term_replaced_by_nn = @. (Rmax - (Rmax - Rmin) * circ / (km + circ))
    lower = minimum(term_replaced_by_nn)
    upper = maximum(term_replaced_by_nn)
    lines!(
        nn_friberg_axis,
        term_replaced_by_nn,
        vec(nn_output) .* k_tr;
        alpha = 0.2,
        linewidth = 3,
        color = :black,
    )
    lines!(
        nn_friberg_axis,
        [lower, upper],
        [lower, upper],
        linestyle = :dash,
        color = :blue,
        linewidth = 2,
    )

end
axislegend(conc_axis, "Dose"; position = :rt, unique = true)
display(fig);

Note here that NN refers to the output of the multi-layer perceptron scaled by k_tr, so that NN replaces the alternative Friberg term,

\[ R_{max} - (R_{max} - R_{min}) * \frac{Circ}{k_m + Circ}.\]

Since the UDE model is being trained on data simulated using the alternative Friberg model, these terms, plotted against each other in the last plot, will ideally agree.

Other diagnostics are done in a similar way as with the true model.

# Figure and axis
figure = Figure(; size = (800, 800), fontsize = 20, title = "UDE")
circ_axis = Axis(figure[1, 1]; title = "Circ", xlabel = "Time")
ylims!(circ_axis, 0.0, nothing)
# Pre-compute one distinct color per subject
subject_colors = cgrad(:oxy, length(ude_eval_subjects); categorical = true)
for (i, (sim_sub, target_sub)) in enumerate(zip(UDE_sims, ude_subjects))
    c = subject_colors[i]
    circ = sim_sub.dynamics.Circ # Extract values from subject
    # Plot circulation trends
    label = sim_sub.subject.events[1].amt |> Int |> string
    lines!(circ_axis, fine_time, circ; color = c, label)
    scatter!(circ_axis, obs_times, target_sub.observations.pd; color = c)
end
axislegend(circ_axis, "Doses"; position = :rb, unique = true)
display(figure)
insp = inspect(fpm_ude)
figure = (; size = (900, 900), fontsize = 20, title = "UDE")
display(convergence_trace(fpm_ude; figure));
display(goodness_of_fit(insp; figure));
display(wresiduals_dist(insp; figure));
display(vpc_plot(ude_model, vpc(fpm_ude; observations = [:pd]); figure));

[ Info: Calculating predictions.
[ Info: Calculating weighted residuals.
[ Info: Calculating empirical bayes.
[ Info: Evaluating dose control parameters.
[ Info: Evaluating individual parameters.
[ Info: Done.
[ Info: Continuous VPC

For the next diagnostic plot, we extract the mean PK and PD trends predicted by the best model, for subjects in the same test set we used in the RNN analysis. The plots look rather similar to the RNN ones:

diagnose_outputs = splits[:test_outputs]
# Get point predictions for PK and PD for each time and each test subject:
error_free_simulated_population = simobs(
    ude_model,
    simulated_population[test_indices],
    coef(fpm_ude);
    simulate_error = false,
)
predictions = zeros(Float32, (2, trend_size, length(test_indices)))
for (subject_ID, subject) in enumerate(error_free_simulated_population)
    # PK trend
    predictions[1, :, subject_ID] .= subject.observations.pk
    # PD trend
    predictions[2, :, subject_ID] .= subject.observations.pd
end
# Plot
fig = Figure(; size = (800, 800), fontsize = 22)
for (col, title) in enumerate(("PK", "PD"))
    # Prepare data for plotting
    targets = @view diagnose_outputs[col, :, :]
    preds = @view predictions[col, :, :]
    plot_info = zip(targets, preds) |> collect |> vec
    start_line = min(minimum(preds), minimum(targets))
    finish_line = max(maximum(preds), maximum(targets))
    # Create Axis and plot
    ylabel = col == 1 ? "Predictions" : ""
    axis = Axis(fig[1, col]; title, ylabel, xlabel = "Observations")
    scatter!(axis, plot_info; marker = 'o')
    lines!(
        axis,
        [(start_line, start_line), (finish_line, finish_line)];
        color = :black,
        linewidth = 2,
    )
end
display(fig);

And here are the PK and PD trend plots for the same test subjects:

diagnose_inputs = splits[:test_inputs]
diagnose_doses = unique(@view diagnose_inputs[2, 1, :])
n_doses = length(diagnose_doses)
# Setup figure and labels
fig = Figure(; size = (900, 170 * (n_doses + 1)), fontsize = 22)
Label(fig[0, 1], "PK"; font = :bold, tellwidth = false)
Label(fig[0, 2], "PD"; font = :bold, tellwidth = false)
# Per dose, plot observations and predictions for PK and PD
row = 0
for dose in diagnose_doses
    row += 1
    # Get subjects and dose value before normalization
    subjects_current_dose = findall(==(dose), @view diagnose_inputs[2, 1, :])
    original_dose = (dose * deviation_dose) + minimum_dose
    dose_label = "Dose " * string(round(Int, original_dose))
    # Create axes for current row
    PK_axis = Axis(fig[row, 1]; ylabel = dose_label)
    ylims!(PK_axis, (0.0, max_targets[1]))
    PD_axis = Axis(fig[row, 2])
    ylims!(PD_axis, (0.0, max_targets[2]))
    # Suppress x axis details in all but the last row
    if row != n_doses
        hidexdecorations!(PK_axis)
        hidexdecorations!(PD_axis)
    else
        PK_axis.xlabel = PD_axis.xlabel = "Time"
    end
    # Iterate in subjects receiving current dose
    for subject_ID in subjects_current_dose
        PK_axis.xgridvisible = PD_axis.xgridvisible = true
        # PK observations
        scatter!(
            PK_axis,
            obs_times,
            @view diagnose_outputs[1, :, subject_ID];
            alpha = 0.5,
            color = :darkgreen,
        )
        # PK RNN predictions
        pk_pred = predictions[1, :, subject_ID]
        lines!(PK_axis, obs_times, pk_pred; linewidth = 2, color = :black)
        # PK error band (±(2*SD) from combined additive-proportional error model)
        pk_std = @. sqrt(σ_add_fit^2 + (pk_pred * σ_prop_fit)^2)
        band!(
            PK_axis,
            obs_times,
            pk_pred .- 2 .* pk_std,
            pk_pred .+ 2 .* pk_std;
            color = (:purple, 0.03),
        )
        # PD observations
        scatter!(
            PD_axis,
            obs_times,
            @view diagnose_outputs[2, :, subject_ID];
            alpha = 0.5,
            color = :darkgreen,
            label = "Observations",
        )
        # PD RNN predictions
        pd_pred = predictions[2, :, subject_ID]
        lines!(
            PD_axis,
            obs_times,
            pd_pred;
            linewidth = 2,
            color = :black,
            label = "Predictions",
        )
        # PD error band (±(2*SD) from additive error model)
        band!(
            PD_axis,
            obs_times,
            pd_pred .- 2 .* σ_add_pd_fit,
            pd_pred .+ 2 .* σ_add_pd_fit;
            color = (:purple, 0.03),
        )
        Label(fig[row, 0], "Test"; font = :bold, tellheight = false, rotation = π / 2)
    end
end
display(fig);

5 Model comparisons

To compare the relative performance of the RNN and UDE models we compute log losses (negative log-likelihoods) on the test subjects in our simulated population. To get a best case scenario baseline, we train the alternative Friberg data generating model on the generated training and validation subjects (used in training and optimizing the RNN and UDE models) and compute its loss on the test subjects as well. We start with that baseline:

out_of_sample_losses = Dict()
# train the data generating model on the training and validation subjects:
fpm_alternative_on_train = fit(
    alternative_model,
    simulated_population[vcat(training_indices, validation_indices)],
    references, # model parameter initialization
    NaivePooled(),
)
# compute negative log likelihood for test subjects:
out_of_sample_losses["Alternative Friberg"] =
    -loglikelihood(
        alternative_model,
        simulated_population[test_indices],
        coef(fpm_alternative_on_train),
        NaivePooled(),
    )

Computing the loss for the optimal UDE model is similar:

out_of_sample_losses["Universal Differential Equation"] =
    -loglikelihood(
        ude_model,
        simulated_population[test_indices],
        coef(fpm_ude),
        NaivePooled(),
    )

For the RNN, we can use the loss function we defined for training purposes:

out_of_sample_losses["RNN"] = loss(best_learned, test_inputs, test_outputs)
out_of_sample_losses

Dict{Any, Any} with 3 entries:
  "RNN"                             => 244.164
  "Universal Differential Equation" => 221.247
  "Alternative Friberg"             => 220.65

As usual, we must interpret these losses carefully, as they are only estimates of the expected loss in the full population, and we do not have associated uncertainty estimates. Still, it appears as if the UDE model performs as well as the RNN model, if not better, despite being trained on much less data. This is to be expected, as the UDE model is still an ODE-based model, and shares some structural similarities with the data-generating alternative Friberg model.

6 Conclusion

This tutorial’s goal is to compare approaches to NLME with varying degrees of data usage. The starting point was a fully mechanistic model, an alternative Friberg model, to create a synthetic population. Then, a fully data-driven RNN model was trained to fit said data.

The point here is that purely ML methods tend to require more data, in this case using 80 subjects for training. Furthermore, these approaches are prone to overfitting, so additional points are required for out-of-sample performance tracking, namely 10 for validation and 10 for testing.

On the other hand, SciML allows for a mix of scientific knowledge and ML. And the researcher has control over the positioning the ML model. This enables fine tuning the scope of the data-driven component, as if the balance between scientific knowledge and data constrains the ML model and fights overfitting, resulting in 8 sufficing for a similar fit.

7 Appendix

7.1 Out-of-sample performance trend

In the following, the function livePlot is defined to occasionally show a plot of the RNN model’s performance in all three splits during training. livePlot can be called from the main() function after each iteration of the hyperparameter optimization loop. Losses for the starting epochs are dropped, suppressing the initial “transient”.

#| output: false
function livePlot(history, hyperparameters)
    # Setup figure and axis
    train, val, test = map(k -> get(history, k)[2], [:train, :val, :test])
    epochs = 1:length(train)
    f = Figure(; size = (700, 700), fontsize = 21)
    yscale = maximum(test) / minimum(test) > 1e2 ? log10 : identity
    ax = Axis(f[1, 1]; yscale, xlabel = "Epochs", ylabel = "Loss")
    # Plots
    plots = Any[map(s -> lines!(ax, epochs, s; linewidth = 2), [train, val, test])...]
    minimum_test_marker =
        scatter!(ax, test |> findmin |> reverse; marker = :x, markersize = 17, color = :red)
    append!(plots, [minimum_test_marker])
    Legend(f[1, 2], plots, ["Training", "Validation", "Test", "Best test"])
    display(f)
    return nothing
end

References

Bartz, Eva, Thomas Bartz-Beielstein, Martin Zaefferer, and Olaf Mersmann. 2023. Hyperparameter Tuning for Machine and Deep Learning with r. Springer.

Friberg, Lena, Anja Henningsson, Hugo Maas, Laurent Nguyen, and Mats Karlsson. 2002. “Model of Chemotherapy-Induced Myelosuppression with Parameter Consistency Across Drugs.” Journal of Clinical Oncology 20 (24): 4713–21. https://doi.org/10.1200/JCO.2002.02.140.

Li, Lisha, Kevin Jamieson, Giulia DeSalvo, Afshin Rostamizadeh, and Ameet Talwalkar. 2018. “Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization.” Journal of Machine Learning Research 18 (185): 1–52. http://jmlr.org/papers/v18/16-558.html.

Soto, Elena, Alexander Staab, Christiane Doege, Matthias Freiwald, Gerd Munzert, and Iñaki F. Trocóniz. 2011. “Comparison of Different Semi-Mechanistic Models for Chemotherapy-Related Neutropenia: Application to BI 2536 a Plk-1 Inhibitor.” Cancer Chemotherapy and Pharmacology 68 (6): 1517–27. https://doi.org/10.1007/s00280-011-1647-3.