diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/README.md b/battery-module-cooling-analysis-with-fourier-neural-operator/README.md new file mode 100644 index 0000000..e093094 --- /dev/null +++ b/battery-module-cooling-analysis-with-fourier-neural-operator/README.md @@ -0,0 +1,151 @@ +# Battery Module Cooling Analysis with Fourier Neural Operator + +This example demonstrates applying a Fourier Neural Operator [1] to heat analysis of a 3D battery module. + +![Plot of true and predicted temperature](./resources/plot.png) +![Plot of difference](./resources/error.png) + +## Setup + +Run the example by running [`example.mlx`](./example.mlx). + +## Requirements + +Requires +- [MATLAB](https://www.mathworks.com/products/matlab.html) (R2025a or newer) +- [Deep Learning Toolbox™](https://www.mathworks.com/products/deep-learning.html) +- [Partial Differential Equation Toolbox™](https://uk.mathworks.com/products/pde.html) +- [Parallel Computing Toolbox™](https://uk.mathworks.com/products/parallel-computing.html) (for training on a GPU) + +## References +[1] Li, Zongyi, Nikola Borislavov Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2021. "Fourier Neural Operator for Parametric Partial Differential Equations." In International Conference on Learning Representations. https://openreview.net/forum?id=c8P9NQVtmnO. + +## Example + +This example applies a 3D Fourier Neural Operator (FNO) to a thermal analysis of a mattery module. The battery module is composed of a number of cells. + +![Image of battery module](./resources/module.png) + +Each cell has the same geometry, as specified by the parameters on this plot. + +![Image of cell](./resources/cell.png) + +The following is an overview of some key parts of the `example.mlx`. + +The [`createBatteryModuleGeometry`](./createBatteryModuleGeometry.m) function creates the battery module geometry as the variable `geomModule`, an instance of [`fegeometry`](https://uk.mathworks.com/help/pde/ug/fegeometry.html). The `geomModule` can be used by the [`femodel`](https://uk.mathworks.com/help/pde/ug/femodel.html) function, which creates an object that can be used to specify and solve the underlying partial differential equation (PDE). + +``` +model = femodel(AnalysisType="thermalTransient", Geometry=geomModule); +``` + +The `"thermalTransient"` input specifies that this is a heat equation that evolves with time. The [`materialProperties`](https://uk.mathworks.com/help/pde/ug/pde.materialproperties.html), [`faceLoad`](https://uk.mathworks.com/help/pde/ug/faceload.html) and [`cellLoad`](https://uk.mathworks.com/help/pde/ug/cellload.html) functions are used to specify the corresponding properties on the domain. In this example the `AmbientTemperature`, `ConvectionCoefficient` and `Heat` for specified faces and cells are set to various values from a collection of potential values `params`, and each instance of the model is solved to generate training data. +``` +results = cell(size(params,1),1); +for i = 1:size(params,1) + model.FaceLoad([boundaryIDs(1).FrontFace, ... + boundaryIDs(end).BackFace]) = ... + faceLoad(ConvectionCoefficient=params(i,2), ... + AmbientTemperature=params(i,1)); + + nominalHeatGen = params(i,3)/volume(1).Cell; + model.CellLoad(cellIDs) = cellLoad(Heat=nominalHeatGen); + + model.CellIC = cellIC(Temperature=params(i,1)); + + T = 60*10; + results{i} = solve(model,0:60:T); +end +``` + +The `results` hold information about the specification and solution of the PDE. These values may be known only at the vertices of the `femodel`-s mesh. However methods like FNO typically require the data to be layed out on a grid, like the grids that `meshgrid` and `ndgrid` produce. The workaround used in this example is to interpolate from the mesh vertices onto a grid using [`scatteredInterpolant`](https://uk.mathworks.com/help/matlab/ref/scatteredinterpolant.html). For example the `ConvectionCoefficient` values that were specified on the front and back face of the battery module are interpolated with the following code. + +``` +F = scatteredInterpolant(result.Mesh.Nodes.', convection); +convectionGrid = F(X,Y,Z); +``` + +This gives us input data `U` and target data `V` on a 32 x 32 x 32 grid. The array `U` is 32 x 32 x 32 x 6 x 216 and `V` is 32 x 32 x 32 x 1 x 216. Each of `U(:,:,:,:,i), V(:,:,:,:,i)` is a distinct observation. The entries `U(:,:,:,1:3,:)` are the $(x,y,z)$ coordinates of the grid, and `U(:,:,:,4:6,:)` are the ambient temperature, convection coefficient and heat generation values respectively, each interpolated from the mesh vertices on which they were specified to the grid. The target `V` is the temperature of the battery module interpolated from the solution on the mesh vertices to the grid points, after the battery has cooled for 10 minutes. + +This data can now be used directly by a 3D FNO. Using the provided [`spectralConvolution3dLayer`](./spectralConvolution3dLayer.m) the following code creates a Fourier layer that combines the spectral convolution with a non-linear activation, a normalization layer and a separate linear layer. +``` +function layer = fourier3dLayer(numModes,hiddenSize,args) +arguments + numModes + hiddenSize + args.Activation = @geluLayer +end +layers = [ + identityLayer(Name="in") + spectralConvolution3dLayer(numModes,hiddenSize) + additionLayer(2,Name="add") + layerNormalizationLayer + args.Activation()]; +net = dlnetwork(layers,Initialize=false); +net = net.addLayers(convolution3dLayer(1,hiddenSize,Name="conv")); +net = connectLayers(net,"conv","add/in2"); +net = connectLayers(net,"in","conv"); +layer = networkLayer(net); +end +``` + +The FNO is completed by connecting a number of Fourier layers in sequence, with linear lifting and projection layers, and additional residual connections. +``` +layers = [ + inputLayer([n,n,n,size(U,4),NaN],'SSSCB',Name="in") + convolution3dLayer(1,2*hiddenSize,Padding="same") + layerNormalizationLayer + geluLayer + convolution3dLayer(1,hiddenSize,Padding="same") + repmat(residualLayer(fourier3dLayer(numModes,hiddenSize)),nlayers,1) + convolution3dLayer(1,2*hiddenSize,Padding="same") + layerNormalizationLayer + geluLayer + convolution3dLayer(1,size(V,4),Padding="same")]; +``` + +This is then trained on the data using [`trainnet`](https://uk.mathworks.com/help/deeplearning/ref/trainnet.html). + +``` +net = dlnetwork(layers,Initialize=false); + +opts = trainingOptions("adam",... + Plots="training-progress",... + MiniBatchSize = 8,... + InputDataFormats="SSSCB",... + TargetDataFormats="SSSCB",... + Shuffle="every-epoch",... + ValidationData = {Uval,Vval},... + ValidationFrequency=100,... + MaxEpochs = 1000); + +net = trainnet(Unorm,Vnorm, net, "l2loss" ,opts); +``` + +The [`minibatchpredict`](https://uk.mathworks.com/help/deeplearning/ref/minibatchpredict.html) function then makes it easy to predict the temperature quickly on the 43 observations in the validation data. This takes about 5s to predict the temperature for all 43 validation observations. +``` +pred = minibatchpredict(net,Uval); +``` + +The predictions are on the 32 x 32 x 32 grid that was specified before. The [`griddedInterpolant`](https://uk.mathworks.com/help/matlab/ref/griddedinterpolant.html) function can interpolate the grid back to the mesh vertices. + +``` +P = [2,1,3]; +result = results{idxVal(i)}; +interpolation = griddedInterpolant(... + permute(X,P),... + permute(Y,P),... + permute(Z,P),... + permute(pred(:,:,:,:,i),P),... + 'spline'); +interpolatedPrediction = interpolation(result.Mesh.Nodes.'); +``` + +Once the data has been interpolated back to the mesh vertices, the [`pdeplot3D`](https://uk.mathworks.com/help/pde/ug/pdeplot3d.html) can be used to create the plots above. + +``` +pdeplot3D(result.Mesh, ColorMapData = interpolatedPrediction, FaceAlpha = 1); +``` + +# + +Copyright 2025 The MathWorks, Inc. diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/createBatteryModuleGeometry.m b/battery-module-cooling-analysis-with-fourier-neural-operator/createBatteryModuleGeometry.m new file mode 100644 index 0000000..3eb1880 --- /dev/null +++ b/battery-module-cooling-analysis-with-fourier-neural-operator/createBatteryModuleGeometry.m @@ -0,0 +1,100 @@ +function [geomModule, domainIDs, boundaryIDs, volume, boundaryArea, ReferencePoint] = createBatteryModuleGeometry(numCellsInModule, cellWidth,cellThickness,tabThickness,tabWidth,cellHeight,tabHeight, connectorHeight ) +%% Uses Boolean geometry functionality in PDE Toolbox, which requires release R2025a or later. +% If you have an older version, use the helper function in this example: +% https://www.mathworks.com/help/pde/ug/battery-module-cooling-analysis-and-reduced-order-thermal-model.html + +% Copyright 2025 The MathWorks, Inc. + +% First, create a single pouch cell by unioning the cell, tab and connector +% Cell creation +cell1 = fegeometry(multicuboid(cellThickness,cellWidth,cellHeight)); +cell1 = translate(cell1,[cellThickness/2,cellWidth/2,0]); +% Tab creation +tab = fegeometry(multicuboid(tabThickness,tabWidth,tabHeight)); +tabLeft = translate(tab,[cellThickness/2,tabWidth,cellHeight]); +tabRight = translate(tab,[cellThickness/2,cellWidth-tabWidth,cellHeight]); +% Union tabs to cells +geomPouch = union(cell1, tabLeft, KeepBoundaries=true); +geomPouch = union(geomPouch, tabRight, KeepBoundaries=true); +% Connector creation +overhang = (cellThickness-tabThickness)/2; +connector = fegeometry(multicuboid(tabThickness+overhang,tabWidth,connectorHeight)); +connectorRight = translate(connector,[cellThickness/2+overhang/2,tabWidth,cellHeight+tabHeight]); +connectorLeft = translate(connector,[(cellThickness/2-overhang/2),cellWidth-tabWidth,cellHeight+tabHeight]); +% Union connectors to tabs +geomPouch = union(geomPouch,connectorLeft,KeepBoundaries=true); +geomPouch = union(geomPouch,connectorRight,KeepBoundaries=true); +% Scale and translate completed pouch cell to create mirrored cell +geomPouchMirrored = translate(scale(geomPouch,[-1 1 1]),[cellThickness,0,0]); +% Union individual pouches to create full module +% Union even-numbered pouch cells together (original cells) +geomForward = fegeometry; +for i = 0:2:numCellsInModule-1 + offset = cellThickness*i; + geom_to_append = translate(geomPouch,[offset,0,0]); + geomForward = union(geomForward,geom_to_append); +end +% Union odd-numbered pouch cells together (mirrored cells) +geomBackward = fegeometry; +for i = 1:2:numCellsInModule-1 + offset = cellThickness*i; + geom_to_append = translate(geomPouchMirrored,[offset,0,0]); + geomBackward = union(geomBackward,geom_to_append); +end +% Union to create completed geometry module +geomModule = union(geomForward,geomBackward,KeepBoundaries=true); +% Rotate and translate the geometry +geomModule = translate(scale(geomModule,[1 -1 1]),[0 cellWidth 0]); +% Mesh the geometry to use query functions for identifying cells and faces +geomModule = generateMesh(geomModule,GeometricOrder="linear"); +% Create Reference Points for each geometry future +ReferencePoint.Cell = [cellThickness/2,cellWidth/2,cellHeight/2]; +ReferencePoint.TabLeft = [cellThickness/2,tabWidth,cellHeight+tabHeight/2]; +ReferencePoint.TabRight = [cellThickness/2,cellWidth-tabWidth,cellHeight+tabHeight/2]; +ReferencePoint.ConnectorLeft = [cellThickness/2,tabWidth,cellHeight+tabHeight+connectorHeight/2]; +ReferencePoint.ConnectorRight = [cellThickness/2,cellWidth-tabWidth,cellHeight+tabHeight+connectorHeight/2]; +% Helper function to get the cell IDs belonging to cell, tab and connector +[~,~,t] = meshToPet(geomModule.Mesh); +elementDomain = t(end,:); +tr = triangulation(geomModule.Mesh.Elements',geomModule.Mesh.Nodes'); +getCellID = @(point,cellNumber) elementDomain(pointLocation(tr,point+(cellNumber(:)-1)*[cellThickness,0,0])); +% Helper function to get the volume of the cells, tabs, and connectors +getVolumeOneCell = @(geomCellID) geomModule.Mesh.volume(findElements(geomModule.Mesh,"region",Cell=geomCellID)); +getVolume = @(geomCellIDs) arrayfun(@(n) getVolumeOneCell(n),geomCellIDs); +% Initialize cell ID and volume structs +domainIDs(1:numCellsInModule) = struct(Cell=[], ... + TabLeft=[],TabRight=[], ... + ConnectorLeft=[],ConnectorRight=[]); +volume(1:numCellsInModule) = struct(Cell=[], ... + TabLeft=[],TabRight=[], ... + ConnectorLeft=[],ConnectorRight=[]); +% Helper function to get the IDs belonging to the left, right, front, back, top and bottom faces +getFaceID = @(offsetVal,offsetDirection,cellNumber) nearestFace(geomModule,... + ReferencePoint.Cell + offsetVal/2 .*offsetDirection ... % offset ref. point to face + + cellThickness*(cellNumber(:)-1)*[1,0,0]); % offset to cell +% Initialize face ID and area structs +boundaryIDs(1:numCellsInModule) = struct(FrontFace=[],BackFace=[], ... + RightFace=[],LeftFace=[], ... + TopFace=[],BottomFace=[]); +boundaryArea(1:numCellsInModule) = struct(FrontFace=[],BackFace=[], ... + RightFace=[],LeftFace=[], ... + TopFace=[],BottomFace=[]); +% Loop over cell, left tab, right tab, left connector, and right connector to get cell IDs and volumes +for part = string(fieldnames(domainIDs))' + partid = num2cell(getCellID(ReferencePoint.(part),1:numCellsInModule)); + [domainIDs.(part)] = partid{:}; + volumesPart = num2cell(getVolume([partid{:}])); + [volume.(part)] = volumesPart{:}; +end +% Loop over front, back, right, left, top, and bottom faces IDs and areas +dimensions = [cellThickness;cellThickness;cellWidth;cellWidth;cellHeight;cellHeight]; +vectors = [-1,0,0;1,0,0;0,1,0;0,-1,0;0,0,1;0,0,-1]; +areaFormula = [cellHeight*cellWidth;cellHeight*cellWidth;cellThickness*cellHeight;cellThickness*cellHeight;cellThickness*cellWidth - tabThickness*tabWidth;cellThickness*cellWidth - tabThickness*tabWidth]; +i = 1; +for face = string(fieldnames(boundaryIDs))' + faceid = num2cell(getFaceID(dimensions(i),vectors(i,:),1:numCellsInModule)); + [boundaryIDs.(face)] = faceid{:}; + areasFace = num2cell(areaFormula(i)*ones(1,numCellsInModule)); + [boundaryArea.(face)] = areasFace{:}; + i = i+1; +end \ No newline at end of file diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx b/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx new file mode 100644 index 0000000..7a1b50c Binary files /dev/null and b/battery-module-cooling-analysis-with-fourier-neural-operator/example.mlx differ diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/resources/cell.png b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/cell.png new file mode 100644 index 0000000..ed87621 Binary files /dev/null and b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/cell.png differ diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/resources/error.png b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/error.png new file mode 100644 index 0000000..99dc3dc Binary files /dev/null and b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/error.png differ diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/resources/module.png b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/module.png new file mode 100644 index 0000000..6aa2382 Binary files /dev/null and b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/module.png differ diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/resources/plot.png b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/plot.png new file mode 100644 index 0000000..ae4a7f0 Binary files /dev/null and b/battery-module-cooling-analysis-with-fourier-neural-operator/resources/plot.png differ diff --git a/battery-module-cooling-analysis-with-fourier-neural-operator/spectralConvolution3dLayer.m b/battery-module-cooling-analysis-with-fourier-neural-operator/spectralConvolution3dLayer.m new file mode 100644 index 0000000..7f64dfc --- /dev/null +++ b/battery-module-cooling-analysis-with-fourier-neural-operator/spectralConvolution3dLayer.m @@ -0,0 +1,90 @@ +classdef spectralConvolution3dLayer < nnet.layer.Layer ... + & nnet.layer.Formattable ... + & nnet.layer.Acceleratable + % spectralConvolution3dLayer A custom layer implementation of + % spectral convolution for data with 3 spatial dimensions. + + % Copyright 2025 The MathWorks, Inc. + + properties + Cin + Cout + NumModes + end + + properties (Learnable) + Weights + end + + methods + function this = spectralConvolution3dLayer(numModes, outChannels, nvargs) + arguments + numModes (1,1) double + outChannels (1,1) double + nvargs.Name {mustBeTextScalar} = "spectralConv3d" + nvargs.Weights = [] + end + + this.Cout = outChannels; + this.NumModes = numModes; + this.Name = nvargs.Name; + this.Weights = nvargs.Weights; + end + + function this = initialize(this, ndl) + inChannels = ndl.Size( finddim(ndl,'C') ); + outChannels = this.Cout; + numModes = this.NumModes; + + if isempty(this.Weights) + this.Cin = inChannels; + this.Weights = 1./(inChannels*outChannels).*( ... + rand([outChannels inChannels numModes numModes numModes]) + ... + 1i.*rand([outChannels inChannels numModes numModes numModes]) ); + end + end + + function y = predict(this, x) + + % Compute the 3d fft and retain only the low frequency modes as + % specified by NumModes. + x = real(x); + x = stripdims(x); + N = size(x, 1); + Nm = this.NumModes; + xft = fft(x, [], 1); + xft = xft(1:Nm,:,:,:,:); + xft = fft(xft, [], 2); + xft = xft(:,1:Nm,:,:,:); + xft = fft(xft, [], 3); + xft = xft(:,:,1:Nm,:,:); + + % Multiply selected Fourier modes with the learnable weights. + xft = permute(xft, [4 5 1 2 3]); + yft = pagemtimes( this.Weights, xft ); + yft = permute(yft, [3, 4, 5, 1, 2]); + + % Make the frequency representation conjugate-symmetric such + % that the inverse Fourier transform is real-valued. + S = floor(N/2)+1 - this.NumModes; + idx = ceil(N/2):-1:2; + yft = cat(1, yft, zeros([S size(yft, 2:5)], 'like', yft)); + yft = cat(1, yft, conj(yft(idx,:,:,:,:))); + + yft = cat(2, yft, zeros([size(yft,1), S, size(yft,3:5)], like=yft)); + yft = cat(2, yft, conj(yft(:,idx,:,:,:))); + + yft = cat(3, yft, zeros([size(yft,[1,2]), S, size(yft,4:5)], like=yft)); + yft = cat(3, yft, conj(yft(:,:,idx,:,:))); + + % Return to physical space via 3d ifft + y = ifft(yft, [], 3, 'symmetric'); + y = ifft(y,[],2, 'symmetric'); + y = ifft(y,[],1, 'symmetric'); + + % Re-apply labels + y = dlarray(y, 'SSSCB'); + y = real(y); + end + end +end \ No newline at end of file