Discovering Dynamic Models from Clinical Data using Neural-PK/PD

James Lu, Genentech

(RowDefault[, (Spacer1), , (Spacer1), ])  #WolframTechConf

Topics

Introduction

Model construction: sequential approach

Neural-PK: predict PharmacoKinetics from dose

Neural-PK/PD: predict PharmacoDynamics in relation to PK

Comparison with state-of-art model

Conclusions

Data: Pharmacokinetics (PK)/Pharmacodynamics (PD)

Drug concentration in plasma (PK), and platelet count (PLT), dose amount (AMT):

Prepend[PKPLTTimeCourseData[[1,2]][[1;;15]], myTableHeadings]//TableForm

Visualization:

plotData[PKPLTTimeCourseData[[1,2]][[1;;15]]]

From Data to Model: Human Mind

Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/HumanModelingParadigm.jpg"]

From Data to Model in the Age of AI: Modeling 2.0

Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/AIModelingParadigm.jpg"]

Grid[List@{BarcodeImage["https://t.co/8ZItYjpQiX", "QR"],Import["NMI_paper.jpg"]}]

Decoder: PK Neural Net

We aim to use a neural net to approximate the vector field, f of the ODE system:

((\[DifferentialD]Y[t])/(\[DifferentialD]t))==Style[f,Blue,Bold][Y[t],α]+(∑_(i=1)^(Subscript[n, d]))(Subscript[dose, i])DiracDelta[t-(Subscript[T, i])]//TraditionalForm

PK Vector Field

Set dimensions of the network:

DataDim=5;
dimPKVar=2;
dimPKParam=4;
timeStepSize=1/4;

Build the neural net:

softPlusLayer=ElementwiseLayer["SoftPlus"];
PKDecoderNetInternal[n_]:= NetInitialize[NetGraph[<| "catenate"->CatenateLayer[1], "lin1"-> LinearLayer[dimPKVar+dimPKParam+n], "SoftPlus1"->softPlusLayer, "lin2"-> LinearLayer[dimPKVar+dimPKParam+n], "SoftPlus2"->softPlusLayer, "lin3"-> LinearLayer[dimPKVar+dimPKParam+n], "SoftPlus3"->softPlusLayer, "lin4"-> LinearLayer[dimPKVar+dimPKParam+n],
"SoftPlus4"->softPlusLayer, "lin5"-> LinearLayer[dimPKVar]|>,{{NetPort["PrevState"], NetPort["PKParam"]}-> "catenate","catenate"->"lin1"->"SoftPlus1"->"lin2"->"SoftPlus2"->"lin3"->"SoftPlus3"->"lin4"->"SoftPlus4"->"lin5"-> NetPort["VFOut" ]},"PrevState"-> {dimPKVar}, "PKParam"->{dimPKParam} , "VFOut"-> {dimPKVar}],"Method"-> {"Random", "Weights"-> NormalDistribution[0.001, 10^-6], "Biases"-> 0}]
PKVFNet=PKDecoderNetInternal[1]

PKVFNet//Information

Decoder: PK Neural Net

Forward Euler Scheme

Consider ODE system of the form:

((\[DifferentialD]Y[t])/(\[DifferentialD]t))==f[Y[t],α]+(∑_(i=1)^(Subscript[n, d]))(Subscript[dose, i])DiracDelta[t-(Subscript[T, i])]//TraditionalForm

With step size Δt, forward Euler time discretization takes the form:

(* Y[t+Δt]==Y[t]+(∫_t^(t+Δt))f[Y[t],α]\[DifferentialD]t +(Subscript[dose, i])UnitStep[t-(Subscript[T, i])] //TraditionalForm *)
Y["t+Δt"]==Y[t]+f[Y[t],α]*Δt+(Subscript[dose, i])UnitStep[t-(Subscript[T, i])]//TraditionalForm

This can be implemented using NetGraph as:

simpleForwEuler=
NetGraph[<| "VFnet"-> PKVFNet, "Δt*"-> ElementwiseLayer[timeStepSize*#&],
"plus"->ThreadingLayer[Plus],"add dose"->ThreadingLayer[Plus],"pad 0"-> PaddingLayer[{{0, dimPKVar-1}}] |>,{NetPort["PKParam"]->NetPort["VFnet", "PKParam"],NetPort["Dose"]-> "pad 0"-> "add dose",NetPort["PrevState"]-> "add dose","add dose"->"VFnet"-> "Δt*"->  "plus"->  NetPort["NewState"],NetPort["PrevState"]-> "plus" },
"PrevState"-> {dimPKVar}, "NewState"-> {dimPKVar},  "PKParam"-> dimPKParam];

Information[simpleForwEuler,"SummaryGraphic"]

Decoder: PK Neural Net

Use NetFoldOperator to map successive state outputs to the state inputs:

NetFoldOperator[simpleForwEuler, {"NewState"-> "PrevState"}, {"PKParam"} ]

Make it into a NetGraph with an initial state:

eulerFoldNet=NetGraph[List@NetFoldOperator[simpleForwEuler, {"NewState"-> "PrevState"}, {"PKParam"} ],{NetPort["InitState"]->NetPort[1,"PrevState"]}]

Test the network with a set of simple input:

eulerFoldNet[<|"InitState"->{0,0},"PKParam"-> {1,1,1,1},
"Dose"-> {{0},{0},{1},{0}}|>]//TableForm

Encoder: PK Neural Net

Pre-processor network: perform nonlinear transformation of data

SELU=ElementwiseLayer["SELU"];
preProcessNet=NetMapThreadOperator@NetChain[{LinearLayer[DataDim], SELU,LinearLayer[DataDim],SELU,LinearLayer[DataDim],SELU,LinearLayer[DataDim] }]

Post-processor network:

postProcessNet[n_]:=NetChain[{LinearLayer[n],SELU,LinearLayer[dimPKParam]}]

postProcessNet[20]

encoderNet[n_]:=NetGraph[<|
"Pre-Process"-> preProcessNet,
"GRU"->GatedRecurrentLayer[n,"Dropout"->0],"last 1"->SequenceLastLayer[], 
"last 2"->SequenceLastLayer[],
"Post-Process"->postProcessNet[n], 
"key"->NetMapOperator[n],
"value"->NetMapOperator[n],
"Attention"->AttentionLayer["Dot"], 
"cat"-> CatenateLayer[]|>,
{"Pre-Process"-> "GRU"-> "last 1",
{"last 1","last 2"}-> "cat", 
"cat"-> "Post-Process", 
"key"->NetPort["Attention","Key"], "value"->NetPort["Attention","Value"], "GRU"-> NetPort["Attention","Query"],
"Attention"-> "last 2"},
"Input"->{"Varying",DataDim-1}];

PKEncoderNet=NetInitialize[encoderNet[20],"Method"-> {"Random", "Weights"-> NormalDistribution[0, 10^-10], "Biases"-> 1} ];

Information[PKEncoderNet,"SummaryGraphic"]

Training PK Net

Implement desired loss function:

lossLayer["weighted MSE"]=ThreadingLayer[
Function[{yObserved,yPredicted,weight},weight*(yObserved-yPredicted)^2]]

extractPKDataNet=NetGraph[<|"columns 1-3"-> NetMapThreadOperator@PartLayer[1;;3], 
"column 5"-> NetMapThreadOperator[PartLayer[5;;-1]], "cat"-> NetMapThreadOperator@CatenateLayer["Inputs"->{{3},{1}} ]|>, {NetPort["DataInput"]-> "columns 1-3", NetPort["DataInput"]->  "column 5", { "columns 1-3",  "column 5" }-> "cat" } , "DataInput"-> {"Varying",DataDim}]

combinedPKNet=NetGraph[<|"PKODE"->eulerFoldNet,"PKPart"-> extractPKDataNet, "PKEncoder"->PKEncoderNet,"loss"->lossLayer["weighted MSE"],"summation"-> SummationLayer[]|>, {NetPort["PKPDData"]-> "PKPart"->"PKEncoder","PKEncoder"-> NetPort["PKODE","PKParam"],  NetPort["PKODE","ObsState"]-> NetPort["ObsState"],{NetPort["Target"],  NetPort["PKODE","ObsState"], NetPort["Weight"]}->"loss"-> "summation"->NetPort["Loss"]},  "Target"->{"Varying",1}];

Information[combinedPKNet,"SummaryGraphic"]

Decoder : Adding PD Neural Net on top of PK

Get the relevant part of the trained PK net:

trainedPKNet=NetExtract[combinedPKNet, {"PKODE",1, "Net", "VFnet"}]

((\[DifferentialD]Y[t])/(\[DifferentialD]t))==Style[f,Blue,Bold][Y[t],α]+(∑_(i=1)^(Subscript[n, d]))(Subscript[dose, i])DiracDelta[t-(Subscript[T, i])]//TraditionalForm

We now build the PD component of the vector field, f:

First, set the dimensions of the PD vector field:

dimPDVar=4;
dimPDParam=10;

PDVFnet[n_]:=Module[{},NetGraph[<|   "catenate"->CatenateLayer[1], "lin1"-> LinearLayer[18+n], "SELU1"-> ElementwiseLayer["SELU"], "lin2"-> LinearLayer[18+n], "SELU2"-> ElementwiseLayer["SELU"], "lin3"-> LinearLayer[18+n], "SELU3"-> ElementwiseLayer["SELU"], "lin4"-> LinearLayer[18+n],
"SELU4"-> ElementwiseLayer["SELU"], 
"lin5"-> LinearLayer[dimPDVar]|>, {{NetPort["PDPrevState"], NetPort["PDParam"],  NetPort["PK:Plasma"]}-> "catenate", "catenate"->"lin1"->"SELU1"->"lin2"->"SELU2"->"lin3"->"SELU3"->"lin4"->"SELU4"->"lin5"-> NetPort["PDVFOut" ]}, "PDPrevState"-> {dimPDVar}, "PDParam"->{dimPDParam} , "PDVFOut"-> {dimPDVar} , "PK:Plasma"-> 1]];

myPDVFnet=PDVFnet[4]

PKPDVFnet=NetGraph[<|"catenate"->CatenateLayer[1], 
"PKVF"-> myPKVFnet, "PDVF"-> NetInitialize[myPDVFnet, "Method"-> {"Random", "Weights"-> NormalDistribution[0, 10^-5], "Biases"-> 0}], "PartPK"-> PartLayer[1;;dimPKVar],"PartPD"-> PartLayer[dimPKVar+1;;dimStateVar]  |>, {NetPort["PrevState"]-> "PartPK"-> NetPort["PKVF", "PrevState"], NetPort["PrevState"]-> "PartPD"-> NetPort["PDVF", "PDPrevState"], {"PKVF", "PDVF"}-> "catenate"-> NetPort["VFOut"]},"PrevState"-> {dimStateVar},"VFOut"-> {dimStateVar}, "PDParam"->{dimPDParam}] ;
Information[PKPDVFnet, "SummaryGraphic"]

Decoder : Adding PD Neural Net on top of PK

dimStateVar=dimPKVar+dimPDVar;
dimParam=dimPKParam+dimPDParam;

eulerStepNet=NetGraph[<|"pad 0"->PaddingLayer[{{0,dimStateVar-1}}],"plus"->ThreadingLayer[Plus],"add dose"->ThreadingLayer[Plus],"VFnet"->PKPDVFnet,"Δt*"->ElementwiseLayer[timeStepSize*#&],"Volume"->PartLayer[2],"Abs"->ElementwiseLayer[Abs[#]&],"PK comp"->PartLayer[1],"divide"->ThreadingLayer@Divide,"reshape1"->ReshapeLayer[{1}],"reshape2"->ReshapeLayer[{1}],"ReLU"->Ramp,"yPD"->PartLayer[dimStateVar],"reshape3"->ReshapeLayer[{1}],"cat"->CatenateLayer[]|>,{NetPort["PrevState"]->"add dose",NetPort["Dose"]->"pad 0"->"add dose","add dose"->"plus","add dose"->"VFnet"->"Δt*"->"plus","plus"->"ReLU"->NetPort["NewState"],NetPort["PrevState"]->"PK comp"->"reshape2",NetPort["PKParam"]->"Volume"->"Abs"->"reshape1",{"reshape2","reshape1"}->"divide",NetPort["PrevState"]->"yPD"->"reshape3",{"divide","reshape3"}->"cat"->NetPort["ObsState"],"divide"->NetPort["VFnet","PK:Plasma"]},"PrevState"->{dimStateVar},"NewState"->{dimStateVar},"PKParam"->dimPKParam];

eulerFoldNet=NetGraph[List@NetFoldOperator[eulerStepNet, {"NewState"-> "PrevState"}, {"PKParam", "PDParam"} ],{NetPort["InitialState"]->NetPort[1,"PrevState"]}];
Information[%,"SummaryGraphic"]

PD Encoders

extractPDDataNet=NetGraph[<|"columns 1-4"->NetMapThreadOperator@PartLayer[1;;4]|>,{NetPort["DataInput"]->"columns 1-4"},"DataInput"->{"Varying",DataDim}]

Initial condition network

ICEncoder[n_]:=NetInitialize[NetGraph[<|
"Pre-Process"-> NetMapThreadOperator@NetChain[{LinearLayer[DataDim-1], ElementwiseLayer["SELU"],LinearLayer[DataDim],ElementwiseLayer["SELU"],LinearLayer[DataDim+2],LinearLayer[DataDim-1] }],
 "RecurrentNet"->GatedRecurrentLayer[n],  "last"->SequenceLastLayer[] |>, {"Pre-Process"-> "RecurrentNet"->"last"},"Input"->{"Varying",DataDim-1}],"Method"-> {"Random", "Weights"-> NormalDistribution[0, 10^-5], "Biases"-> 0} ];

ICNet=NetInitialize[NetGraph[<| "PDPart"-> extractPDDataNet,"PDEncoder"->NetBidirectionalOperator[ICEncoder[5]], "linear"-> LinearLayer[1],"rep"-> ReplicateLayer[dimPDVar-1],"flatten"-> FlattenLayer[], "pad 0x3"-> PaddingLayer[{{dimPKVar+1,0}}] |>, {NetPort["PKPDData"]->"PDPart"-> "PDEncoder"->  "linear"-> "rep"-> "flatten"-> "pad 0x3"}],"Method"-> {"Random", "Weights"-> NormalDistribution[0, 10^-5], "Biases"-> 0.1} ];
Information[%, "SummaryGraphic"]

PD Parameter Network

RNNNetForward=GatedRecurrentLayer[50,"Input"->{"Varying",DataDim-1}];
PDParamEncoder[n_]:=NetInitialize[NetGraph[<|
"Pre-Process"-> NetMapThreadOperator@NetChain[{LinearLayer[DataDim-1], ElementwiseLayer["SELU"],LinearLayer[DataDim+2],ElementwiseLayer["SELU"],LinearLayer[DataDim+2],ElementwiseLayer["SELU"],LinearLayer[DataDim+2],ElementwiseLayer["SELU"],LinearLayer[DataDim-1] }],
 "GRU"->RNNNetForward,  "last"->SequenceLastLayer[],"Post-Process"->NetChain[{LinearLayer[n],ElementwiseLayer["SELU"],DropoutLayer[0],LinearLayer[dimPDParam]}]|>, {"Pre-Process"-> "GRU"->"last"->"Post-Process"->NetPort["ParamEst"]},"Input"->{"Varying",DataDim-1}],"Method"-> {"Random", "Weights"-> NormalDistribution[0, 10^-5], "Biases"-> 0} ];

Information[PDParamEncoder[50],"SummaryGraphic"]

Complete Model

lossLayer["Standard"]=ThreadingLayer[
Function[{yObserved,yPredicted,weight},weight*(yPredicted-yObserved)^2]];
combinedPKPDNet=NetGraph[<|"ODE"->ODEdecoder, "PKEncoder"->PKencoderNet,"PDEncoder"->PDParamEncoder[50],"loss"->lossLayer["Standard"],"summation"-> SummationLayer[], "PKData"-> extractPKDataNet, "PDData"-> extractPDDataNet, "ICNet"-> ICNet, "Sum"-> TotalLayer[],"PD obs"->NetChain[{NetMapThreadOperator@PartLayer[2], NetMapThreadOperator@ReshapeLayer[{1}]}]|>, {NetPort["PKPDData"]-> "PKData"->"PKEncoder",NetPort["PKPDData"]->"PDData"-> "PDEncoder", "PKEncoder"-> NetPort["ODE","PKParam"],   "PDEncoder"-> NetPort["ODE","PDParam"],NetPort["ODE","ObsState"]-> NetPort["ObsState"],NetPort["ODE","ObsState"]-> "PD obs",{NetPort["Target"], "PD obs", NetPort["Weight"]}->"loss"-> "summation", NetPort["PKPDData"]-> "ICNet", {"ICNet",NetPort["InitialState"]}-> "Sum","Sum"-> NetPort["ODE","InitialState"],"summation"->NetPort["Loss"]},  "Target"->{"Varying",1},"Weight"->{"Varying",1}];
Information[%,"SummaryGraphic"]

NetTrain[combinedPKPDNet, "...", LearningRateMultipliers->{"PKEncoder"-> 0, {"ODE",1, "Net", "VFnet", "PKVF" }-> 0 , _-> 1}]

Clinical Data

Compiled data: longitudinal platelet response from 665 patients treated with T-DM1

Phase I: various doses and schedules

3×Phase II, 1×Phase III: 3.6 mg/kg, Q3W

Train and test split

Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/TrainTestSplit.jpg"]

Neural-PK/PD Recapitulates Qualitative Dynamics

Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/NeuralPKPD_QualitativeDynamics.jpg"]

See it in action:

SetDirectory[NotebookDirectory[]];
stateTrajVideo=Import["Code/myPKPDStateTraj.mp4"]

Neural-PK/PD Enables Accurate Individual Predictions

Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/NeuralPKPD_Predictions.jpg"]

Demonstration on Clinical Data

Import["Code/myPKPDSim.mp4"]

Neural-PK/PD Simulation of Dose Schedules

Grid@List@{Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/Simulation_Q3W_Q1W.jpg"], Import["/Users/luj25/Google Drive/Projects/MachineLearning/Presentations/WTC_2021/Code/Simulation_Q3D.jpg"]}

Conclusions

Wolfram Language provides high level functionalities enabling construction of complex neural nets for time course data with ease:

NetGraph, NetFoldOperator, ...

Compared to other DL frameworks, the construction of PK and PK/PD models in a sequential manner is eased by the ability to extract/combine models and train only subparts thereof:

NetExtract

NetTrain[..., LearningRateMultipliers->{ ...-> 0 , _-> 1}]

Future development for scientists/engineers who wish to apply DL to dynamical systems: can there be a NeuralODELayer based on various ODE solvers?

Acknowledgement

Genentech: Brendan Bender, Gengbo Liu, Rashed Harun, Omid Bazigir, Dan Lu, Kai Liu, Jin Jin, CPLT & Amita Joshi

University of Michigan: Prof. Yuanfang Guan, Kaiwen Deng

Ann Arbor Algorithms: Xinyuan Zhang

Wolfram Research: Shay Shepston, Wolfram Events Team