diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx b/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx
index 7a1b50c..47d3b80 100644
Binary files a/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx and b/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx differ
diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/trainingPartitions.m b/battery-module-cooling-analysis-with-fourier-neural-operator/trainingPartitions.m
new file mode 100644
index 0000000..2f2a39f
--- /dev/null
+++ b/battery-module-cooling-analysis-with-fourier-neural-operator/trainingPartitions.m
@@ -0,0 +1,47 @@
+function varargout = trainingPartitions(numObservations,splits)
+%TRAININGPARTITONS Random indices for splitting training data
+%   [idx1,...,idxN] = trainingPartitions(numObservations,splits) returns
+%   random vectors of indices to help split a data set with the specified
+%   number of observations, where SPLITS is a vector of length N of
+%   partition sizes that sum to one.
+%
+%   % Example: Get indices for 50%-50% training-test split of 500
+%   % observations.
+%   [idxTrain,idxTest] = trainingPartitions(500,[0.5 0.5])
+%
+%   % Example: Get indices for 80%-10%-10% training, validation, test split
+%   % of 500 observations. 
+%   [idxTrain,idxValidation,idxTest] = trainingPartitions(500,[0.8 0.1 0.1])
+
+arguments
+    numObservations (1,1) {mustBePositive}
+    splits {mustBeVector,mustBeInRange(splits,0,1,"exclusive"),mustSumToOne}
+end
+
+numPartitions = numel(splits);
+varargout = cell(1,numPartitions);
+
+idx = randperm(numObservations);
+
+idxEnd = 0;
+
+for i = 1:numPartitions-1
+    idxStart = idxEnd + 1;
+    idxEnd = idxStart + floor(splits(i)*numObservations) - 1;
+
+    varargout{i} = idx(idxStart:idxEnd);
+end
+
+% Last partition.
+varargout{end} = idx(idxEnd+1:end);
+
+end
+
+function mustSumToOne(v)
+% Validate that value sums to one.
+
+if sum(v,"all") ~= 1
+    error("Value must sum to one.")
+end
+
+end