%% Apply Neural ODE Operation
% For the initial conditions, create a formatted dlarray object containing a batch of 128 28-by-28 images with 64 channels. Specify the format "SSCB" (spatial, spatial, channel, batch).

miniBatchSize = 128;
inputSize = [28 28];
numChannels = 64;
Y0 = rand(inputSize(1),inputSize(2),numChannels,miniBatchSize);
dlY0 = dlarray(Y0,"SSCB");

% View the size and format of the initial conditions.

size(dlY0)
dims(dlY0)

% Specify the ODE function. Define the function odeModel, listed in the ODE Function section of the example, which applies a convolution operation followed by a hyperbolic tangent operation to the input data.

odefun = @odeModel;

% Initialize the parameters for the convolution operation in the ODE function. The output size of the ODE function must match the size of the initial conditions, so specify the same number of filters as the number of input channels.

filterSize = [3 3];
numFilters = numChannels;

parameters = struct;
parameters.Weights = dlarray(rand(filterSize(1),filterSize(2),numChannels,numFilters));
parameters.Bias = dlarray(zeros(1,numFilters));

% Specify an interval of integration of [0 0.1].

tspan = [0 0.1];

% Apply the neural ODE operation.

dlY = dlode45(odefun,tspan,dlY0,parameters);

% View the size and format of the output.

size(dlY)
dims(dlY)

%
%% ODE Function
% The ODE function odeModel takes as input the function inputs t (unused) and y, and the ODE function parameters p containing the convolution weights and biases, and returns the output of the convolution-tanh block operation. The convolution operation applies padding such that the output size matches the input size.

function z = odeModel(t,y,p)

weights = p.Weights;
bias = p.Bias;
z = dlconv(y,weights,bias,Padding="same");
z = tanh(z);

end

% Copyright 2021 The MathWorks, Inc.