Message Boards

GROUPS:

Staff Picks Biological Sciences Data Science Mathematics Calculus Equation Solving Wolfram Language Machine Learning Publication Materials Neural Networks

James Lu

Advancing towards Modeling 2.0 with neural-ODE via Wolfram Language

James Lu, Genentech, Inc.

Posted 7 months ago

4300 Views

2 Replies

16 Total Likes

Follow this post

SUPPLEMENTARY WOLFRAM MATERIALS for the ARTICLE:

Lu, J., Bender, B., Jin, J.Y. et al. (2021).
Deep learning prediction of patient response time course from early data via neural-pharmacokinetic/pharmacodynamic modelling.
NATURE Machine Intelligence 3, 696–704.
https://doi.org/10.1038/s42256-021-00357-4
Full article PDF

Neural-PK/PD simulation in action
Video
Neural-PK/PD prediction update in action
Video

Advancing towards Modeling 2.0 with neural-ODE via Wolfram Language
by James Lu
Principal AI Scientist, Genentech

Introduction

Since the days of Newton, scientists have used differential equations to describe the physical world. Alan Turing also once wrote on a postcard, “Science is a Differential Equation” (from Hodges, A., 2014):

In[]:=

How have we gone about the process of creating differential equation models of data? In the field of Pharmacometrics modeling of clinical data, perhaps the following figure summarizes it :

In[]:=

However, with the advent of Deep Learning and the notion of Software 2.0, does that open up a new paradigm? We proposed the following as Modeling 2.0:

In[]:=

We’ve worked out how to do this in the Pharmacokinetics/Pharmacodynamics (PK/PD) setting (based on the seminal work of [Chen et al, NeurIPS 2018]) and proposed Neural-PK/PD methodology which was published in Nature Machine Intelligence (2021):

In the following, we give an overview of the data involved and how to build the PK and PD submodules respectively in a sequential manner. You’ll see how the basic layers offered by Wolfram Language is expressively enough to implement this formalism. Finally, we will show how to visualize the inner workings of such models.

Data

Here, we set some hyperparameters of the model: in particular, our data will have 5 columns (“TAD (day)”,”TIME (day)”,”PK”,”PD”, “AMT(dosed)”):

In[]:=

DataDim=5;timeStepSize=1/4;

In[]:=

trainPKPDData={0,0,6.8146`,-100,222.`},{1.`,1,52.298`,196.`,-100},{8.`,8,17.141`,126.`,-100},

;

Here’s a visualization of the data:

Out[]=

Hyperparameters & Data Processing

We set the PK and PD variable and parameter dimensions: here, we can leverage prior intuition from ODE modeling regarding the complexity of the model required -

In[]:=

dimPKVar=2;dimPKParam=4;dimPDVar=4;dimPDParam=10;

Let’s compute the data normalization function:

In[]:=

feTimeCourse=FeatureExtraction[trainPKPDData]

Out[]=

FeatureExtractorFunction

Input type:

Mixed

(number: 5)

Output type: NumericalVector

(length: 5)



Verify the mean and std of transformed data:

In[]:=

{Mean@#,StandardDeviation@#}&/@Transpose[feTimeCourse[trainPKPDData]]//Chop//TableForm

Out[]//TableForm=

0	1.01242
0	1.01242
0	1.01242
0	1.01242
0	1.01242

We then transform the training data:

In[]:=

inputData["Train"]=feTimeCourse/@{trainPKPDData};

Pharmaco-Kinetics Submodule

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

Out[]//TraditionalForm=

Y(t)

t



∑

i=1

dose

δ(t-

)+f(Y(t),α)

Vector field

In[]:=

PKDecoderNetInternal[n_]:=NetInitialize[NetGraph[<|"catenate"CatenateLayer[1],"lin1"LinearLayer[dimPKVar+dimPKParam+n],"SoftPlus1"ElementwiseLayer["SoftPlus"],"lin2"LinearLayer[dimPKVar+dimPKParam+n],"SoftPlus2"ElementwiseLayer["SoftPlus"],"lin3"LinearLayer[dimPKVar+dimPKParam+n],"SoftPlus3"ElementwiseLayer["SoftPlus"],"lin4"LinearLayer[dimPKVar+dimPKParam+n],"SoftPlus4"ElementwiseLayer["SoftPlus"],"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,10^-10],"Biases"0}]

In[]:=

PKVFNet=PKDecoderNetInternal[1]

Out[]=

NetGraph



Number of inputs:	2
VFOut port:	vector (size: 2)



We can check the initialized network:

In[]:=

PKVFNet[<|"PrevState"{0,0},"PKParam"{1,1,1,1}|>]

Out[]=

{-4.29992×

-11

,-1.96244×

-11

}

Euler step

This performs the computation required for 1 Euler step:

In[]:=

eulerStepNet=NetGraph[<|"pad 0"PaddingLayer[{{0,dimPKVar-1}}],"plus"ThreadingLayer[Plus],"add dose"ThreadingLayer[Plus],"VFnet"PKVFNet,"Δt*"ElementwiseLayer[timeStepSize*#&],"Volume"PartLayer[2],"Abs"ElementwiseLayer[Abs[#]&],"Obs comp"PartLayer[1],"divide"ThreadingLayer@Divide,"reshape1"ReshapeLayer[{1}],"reshape2"ReshapeLayer[{1}],"ReLU"Ramp|>,{NetPort["PrevState"]"add dose",NetPort["Dose"]"pad 0""add dose","add dose""plus","add dose""VFnet""Δt*""plus","plus""ReLU"NetPort["NewState"],NetPort["PrevState"]"Obs comp""reshape2",NetPort["PKParam"]"Volume""Abs""reshape1",{"reshape2","reshape1"}"divide"NetPort["ObsState"]},"PrevState"{dimPKVar},"NewState"{dimPKVar},"PKParam"dimPKParam]

Out[]=

NetGraph



Number of inputs:	3
Number of outputs:	2



However, we need to perform many steps to reach the end of the time domain, which we do via the following:

In[]:=

eulerFoldNet=NetGraph[List@NetFoldOperator[eulerStepNet,{"NewState""PrevState"},{"PKParam"}],{NetPort["InitialState"]NetPort[1,"PrevState"]}]

Out[]=

NetGraph



Number of inputs:	3
Number of outputs:	2



PK Time-Course Encoder

In[]:=

encoderNet[n_,dropout_]:=NetGraph[<|"Pre-Process"NetMapThreadOperator@NetChain[{LinearLayer[DataDim],ElementwiseLayer["SELU"],LinearLayer[DataDim],ElementwiseLayer["SELU"],LinearLayer[DataDim],ElementwiseLayer["SELU"],LinearLayer[DataDim]}],"GRU"GatedRecurrentLayer[n,"Dropout"dropout],"last 1"SequenceLastLayer[],"last 2"SequenceLastLayer[],"Post-Process"NetChain[{LinearLayer[n],ElementwiseLayer["SELU"],DropoutLayer[],LinearLayer[dimPKParam]}],"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}];

In[]:=

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

Out[]=

NetGraph



InputPort
Input:	matrix (size: n ×4)
OutputPort
Output:	vector (size: 4)



Loss function

In[]:=

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

Combined net

Finally, putting all the pieces together for our Neural-PK model:

In[]:=

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}];

In[]:=

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},"Weight"->{"Varying",1}]

Out[]=

NetGraph



Number of inputs:	5
Number of outputs:	3



Train

Let’s try training it on just 1 single patient, replicated 10 times for the sake of example:

In[]:=

n=10;trainedPK=NetTrain[combinedPKNet,<|"PKPDData"Table[trainPKPDData,n],"Target"Table[inputPKObs,n],"Weight"Table[inputWeight,n],"InitialState"Table[initState,n],"Dose"Table[inputDosing,n]|>,LossFunction"Loss",ValidationSetNone,MaxTrainingRounds10,Method{"ADAM","LearningRate"10^-4,"GradientClipping"1,"L2Regularization"0.1}]

Pharmaco-Dynamics Submodule

Having trained and saved the Neural-PK model, one can just load it and take out the part that we want to use further:

In[]:=

pretrainPKNet=Import["Results/NeuralPK_Attention_Final.wlnet"];

In[]:=

myPKVFnet=NetExtract[pretrainPKNet,{"PKODE",1,"Net","VFnet"}]

Out[]=

NetGraph



InputPorts
PrevState:	vector (size: 2)
PKParam:	vector (size: 4)
OutputPort
VFOut:	vector (size: 2)



In[]:=

PKencoderNet=NetExtract[pretrainPKNet,"PKEncoder"];

Vector Field

We now construct a combined Vector Field net for both PK and PD:

In[]:=

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

In[]:=

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["CP"]}"catenate","catenate"->"lin1""SELU1"->"lin2""SELU2"->"lin3""SELU3"->"lin4""SELU4"->"lin5"NetPort["PDVFOut"]},"PDPrevState"{dimPDVar},"PDParam"{dimPDParam},"PDVFOut"{dimPDVar},"CP"1]];

In[]:=

myPDVFnet=NetInsertSharedArrays[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}]

Out[]=

NetGraph



InputPorts
PrevState:	vector (size: 6)
PKParam:	vector (size: 4)
PDParam:	vector (size: 10)
CP:	vector (size: 1)
OutputPort
VFOut:	vector (size: 6)



Euler Step

A pretty complex looking network model can be built up compactly using Wolfram Language:

In[]:=

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","CP"]},"PrevState"{dimStateVar},"NewState"{dimStateVar},"PKParam"dimPKParam]

Out[]=

NetGraph



InputPorts
PrevState:	vector (size: 6)
Dose:	vector (size: 1)
PKParam:	vector (size: 4)
PDParam:	vector (size: 10)
OutputPorts
NewState:	vector (size: 6)
ObsState:	vector (size: 2)



In[]:=

eulerFoldNet=NetGraph[List@NetFoldOperator[eulerStepNet,{"NewState""PrevState"},{"PKParam","PDParam"}],{NetPort["InitialState"]NetPort[1,"PrevState"]}];

PD Encoders

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

Out[]=

NetGraph



InputPort
PKPDData:	matrix (size: n ×5)
OutputPort
Output:	vector (size: 6)



In[]:=

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}]

Combined Net

In[]:=

lossLayer["Standard"]=ThreadingLayer[Function[{yObserved,yPredicted,weight},weight*(yPredicted-yObserved)^2]];combinedPKPDNet=NetGraph[<|"ODE"eulerFoldNet,"PKEncoder"PKencoderNet,"PDEncoder"PDPa