diff --git a/+kwave/+tests/+unit/TestFourierCollocation.m b/+kwave/+tests/+unit/TestFourierCollocation.m index 6a9ba65d..feff03ce 100644 --- a/+kwave/+tests/+unit/TestFourierCollocation.m +++ b/+kwave/+tests/+unit/TestFourierCollocation.m @@ -87,6 +87,38 @@ function testDivergence(testCase) end + % Test the gradient of a vector function. + function testGradientVector(testCase) + import matlab.unittest.constraints.IsEqualTo + + % No staggering. + [f, testCase.referenceSolution] = testCase.getPeriodicTensorFunction; + testCase.actualSolution = testCase.solver.gradientVector(f); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Forward staggering. + [f, testCase.referenceSolution] = testCase.getPeriodicTensorFunction("forward"); + testCase.actualSolution = testCase.solver.gradientVector(f, Staggering="forward"); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Backward staggering. + [f, testCase.referenceSolution] = testCase.getPeriodicTensorFunction("backward"); + testCase.actualSolution = testCase.solver.gradientVector(f, Staggering="backward"); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Scalar kappa. + testCase.solver.kappa = 2; + [f, testCase.referenceSolution] = testCase.getPeriodicTensorFunction; + testCase.referenceSolution = testCase.referenceSolution .* testCase.solver.kappa; + testCase.actualSolution = testCase.solver.gradientVector(f); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Test incorrect size gives exception. + f = rand(2, 2, 2, 4); + testCase.verifyError(@() testCase.solver.gradientVector(f), 'FourierCollocation:incorrectSize'); + + end + % Test the curl function. function testCurl(testCase) import matlab.unittest.constraints.IsEqualTo @@ -303,17 +335,17 @@ function testStagger(testCase) switch staggering case 'none' - xSg = obj.kgridPadded.xVec; - ySg = obj.kgridPadded.yVec; - zSg = obj.kgridPadded.zVec; + xSg = obj.kgridPadded.xVec; + ySg = obj.kgridPadded.yVec; + zSg = obj.kgridPadded.zVec; case 'forward' - xSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; - ySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; - zSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + xSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + ySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + zSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; case 'backward' - xSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; - ySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; - zSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + xSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + ySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + zSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; end switch obj.kgridPadded.dimensions @@ -327,9 +359,9 @@ function testStagger(testCase) [X, Y] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec); [Xsg, Ysg] = ndgrid(xSg, ySg); - + F = sin(kx .* X) .* sin(ky .* Y) ./ kx; - + gradF = zeros([size(F), 1, 2]); gradF(:, :, :, 1) = cos(kx .* Xsg) .* sin(ky .* Y); gradF(:, :, :, 2) = sin(kx .* X) .* cos(ky .* Ysg) .* (ky ./ kx); @@ -337,12 +369,12 @@ function testStagger(testCase) kx = (2*pi ./ obj.kgridPadded.xSize); ky = (2*pi ./ obj.kgridPadded.ySize); kz = (2*pi ./ obj.kgridPadded.zSize); - + [X, Y, Z] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec, obj.kgridPadded.zVec); [Xsg, Ysg, Zsg] = ndgrid(xSg, ySg, zSg); - + F = sin(kx .* X) .* sin(ky .* Y) .* sin(kz .* Z) ./ kx; - + gradF = zeros([size(F), 3]); gradF(:, :, :, 1) = cos(kx .* Xsg) .* sin(ky .* Y) .* sin(kz .* Z); gradF(:, :, :, 2) = sin(kx .* X) .* cos(ky .* Ysg) .* sin(kz .* Z) .* (ky ./ kx); @@ -363,17 +395,17 @@ function testStagger(testCase) switch staggering case 'none' - xSg = obj.kgridPadded.xVec; - ySg = obj.kgridPadded.yVec; - zSg = obj.kgridPadded.zVec; + xSg = obj.kgridPadded.xVec; + ySg = obj.kgridPadded.yVec; + zSg = obj.kgridPadded.zVec; case 'forward' - xSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; - ySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; - zSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + xSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + ySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + zSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; case 'backward' - xSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; - ySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; - zSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + xSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + ySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + zSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; end switch obj.kgridPadded.dimensions @@ -384,30 +416,30 @@ function testStagger(testCase) case 2 kx = (2*pi ./ obj.kgridPadded.xSize); ky = (2*pi ./ obj.kgridPadded.ySize); - + [X, Y] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec); [Xsg, Ysg] = ndgrid(xSg, ySg); - + Fx = sin(kx .* X) ./ kx; Fy = sin(ky .* Y) ./ ky; - + F = cat(4, Fx, Fy); - + divF = cos(kx .* Xsg) + cos(ky .* Ysg); case 3 kx = (2*pi ./ obj.kgridPadded.xSize); ky = (2*pi ./ obj.kgridPadded.ySize); kz = (2*pi ./ obj.kgridPadded.zSize); - + [X, Y, Z] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec, obj.kgridPadded.zVec); [Xsg, Ysg, Zsg] = ndgrid(xSg, ySg, zSg); - + Fx = sin(kx .* X) ./ kx; Fy = sin(ky .* Y) ./ ky; Fz = sin(kz .* Z) ./ kz; - + F = cat(4, Fx, Fy, Fz); - + divF = cos(kx .* Xsg) + cos(ky .* Ysg) + cos(kz .* Zsg); end end @@ -633,6 +665,109 @@ function testStagger(testCase) curlF(:,:,:,3) = dFydx - dFxdy; end + + + % Define a periodic vector function and its analytic gradient on + % the grid specified by obj.kgridPadded, returning the gradients + % in each axis as a tensor field. The function is normalized so + % the maximum of the gradient in each axis is approximately 1. + % The tensor field can also be returned on a staggered grid. + function [F, gradF] = getPeriodicTensorFunction(obj, staggering) + + arguments + obj + staggering(1,:) char {mustBeMember(staggering, {'none', 'forward', 'backward'})} = 'none' + end + + switch staggering + case 'none' + xxSg = obj.kgridPadded.xVec; + xySg = obj.kgridPadded.yVec; + xzSg = obj.kgridPadded.zVec; + yxSg = obj.kgridPadded.xVec; + yySg = obj.kgridPadded.yVec; + yzSg = obj.kgridPadded.zVec; + zxSg = obj.kgridPadded.xVec; + zySg = obj.kgridPadded.yVec; + zzSg = obj.kgridPadded.zVec; + case 'forward' + xxSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + xySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + xzSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + yxSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + yySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + yzSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + zxSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + zySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + zzSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + + case 'backward' + xxSg = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + xySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + xzSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + yxSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + yySg = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + yzSg = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + zxSg = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + zySg = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + zzSg = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + end + + switch obj.kgridPadded.dimensions + case 1 + kx = (2*pi ./ obj.kgridPadded.xSize); + F = sin(kx .* obj.kgridPadded.xVec) ./ kx; + gradF = cos(kx .* xxSg); + case 2 + kx = (2*pi ./ obj.kgridPadded.xSize); + ky = (2*pi ./ obj.kgridPadded.ySize); + + [X, Y] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec); + [xxsg, xysg] = ndgrid(xxSg, xySg); + [yxsg, yysg] = ndgrid(yxSg, yySg); + + Fx = sin(kx .* X) .* sin(ky .* Y) ./ kx; + Fy = sin(ky .* Y) .* sin(kx .* X) ./ ky; + + gradFx_x = cos(kx .* xxsg) .* sin(ky .* Y); + gradFx_y = sin(kx .* X) .* cos(ky .* xysg) .* (ky ./ kx); + gradFy_x = sin(ky .* Y) .* cos(kx .* yxsg) .* (kx ./ ky); + gradFy_y = cos(ky .* yysg) .* sin(kx .* X); + + F = cat(4, Fx, Fy); + gradF = cat(5, cat(4, gradFx_x, gradFy_x), cat(4, gradFx_y, gradFy_y)); + + case 3 + kx = (2*pi ./ obj.kgridPadded.xSize); + ky = (2*pi ./ obj.kgridPadded.ySize); + kz = (2*pi ./ obj.kgridPadded.zSize); + + [X, Y, Z] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec, obj.kgridPadded.zVec); + [xxsg, xysg, xzsg] = ndgrid(xxSg, xySg, xzSg); + [yxsg, yysg, yzsg] = ndgrid(yxSg, yySg, yzSg); + [zxsg, zysg, zzsg] = ndgrid(zxSg, zySg, zzSg); + + Fx = sin(kx .* X) .* sin(ky .* Y) .* sin(kz .* Z) ./ kx; + Fy = sin(ky .* Y) .* sin(kx .* X) .* sin(kz .* Z) ./ ky; + Fz = sin(kz .* Z) .* sin(kx .* X) .* sin(ky .* Y) ./ kz; + + gradFx_x = cos(kx .* xxsg) .* sin(ky .* Y) .* sin(kz .* Z); + gradFx_y = sin(kx .* X) .* cos(ky .* xysg) .* sin(kz .* Z) .* (ky ./ kx); + gradFx_z = sin(kx .* X) .* sin(ky .* Y) .* cos(kz .* xzsg) .* (kz ./ kx); + + gradFy_x = sin(ky .* Y) .* cos(kx .* yxsg) .* sin(kz .* Z) .* (kx ./ ky); + gradFy_y = cos(ky .* yysg) .* sin(kx .* X) .* sin(kz .* Z); + gradFy_z = sin(ky .* Y) .* sin(kx .* X) .* cos(kz .* yzsg) .* (kz ./ ky); + + gradFz_x = sin(kz .* Z) .* cos(kx .* zxsg) .* sin(ky .* Y) .* (kx ./ kz); + gradFz_y = sin(kz .* Z) .* sin(kx .* X) .* cos(ky .* zysg) .* (ky ./ kz); + gradFz_z = cos(kz .* zzsg) .* sin(kx .* X) .* sin(ky .* Y); + + F = cat(4, Fx, Fy, Fz); + gradF = cat(5, cat(4, gradFx_x, gradFy_x, gradFz_x), cat(4, gradFx_y, gradFy_y, gradFz_y), cat(4, gradFx_z, gradFy_z, gradFz_z)); + end + end + end end diff --git a/+kwave/+toolbox/@FourierCollocation/gradientVector.m b/+kwave/+toolbox/@FourierCollocation/gradientVector.m new file mode 100644 index 00000000..c9f4b451 --- /dev/null +++ b/+kwave/+toolbox/@FourierCollocation/gradientVector.m @@ -0,0 +1,161 @@ +%% gradient +% *Class:* kwave.toolbox.FourierCollocation +% *Package:* kwave.toolbox +% +% Calculate gradient of vector field. +% +%% Syntax +% df = gradientVector(obj, f) +% df = gradientVector(obj, f, Staggering='forward') +% +%% Description +% Calculates the gradient of a vector field in 1D, 2D, or 3D using a +% Fourier collocation spectral method. +% +% The tensor components of the gradient are stacked in the 4th and 5th +% dimensions of the output. For example, if calling gradient on a matrix +% of dimensions (10, 10), the output will be of size (10, 10, 1, 2, 2). +% This is to allow codes to implement multi-dimensional support by always +% looping over the fourth and fifth dimensions. +% +% The resulting tensor takes the shape +% df = [dfxdx dfxdy dfxdz; +% dfydx dfydy dfydz; +% dfzdx dfzdy dfzdz] +% +% If obj.kappa is defined, a k-space correction is applied as part of the +% gradient calculation. If kappa is a scalar (single frequency correction) +% or empty, the gradient components are calculated using 1D FFTs. If kappa +% is a matrix, the gradient components are calculated using ND FFTs, and +% kappa is applied in the Fourier domain. +% +% The gradient operations (and kappa if defined) are defined on the padded +% grid. Thus, the inputs to this function must also be defined on the +% padded grid. The output can be returned on a spatially staggered grid by +% setting the optional |Staggering| argument. +% +%% Input Arguments +% * |f| - (numeric) Vector field to compute gradient of. +% +%% Name-Value Arguments +% Specify optional pairs of arguments as |Name1=Value1,...,NameN=ValueN|, +% where |Name| is the argument name and |Value| is the corresponding value. +% Name-value arguments must appear after other arguments, but the order of +% the pairs does not matter. +% +% * |Staggering| - ('none', 'forward', 'backward') Option to return the +% output staggered by half a grid point in the specified direction. +% Default = 'none'. +% +%% Output Arguments +% * |df| - (numeric) Gradient of f. + +% Copyright (C) 2024- University College London. +% +% This file is part of k-Wave-II (http://www.k-wave.org). k-Wave-II is free +% software: you can redistribute it and/or modify it under the terms of the +% GNU Lesser General Public License as published by the Free Software +% Foundation, either version 3 of the License, or (at your option) any +% later version. +% +% k-Wave-II is distributed in the hope that it will be useful, but WITHOUT +% ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +% FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public +% License for more details. +% +% You should have received a copy of the GNU Lesser General Public License +% along with k-Wave-II. If not, see . + +function df = gradientVector(obj, f, options) + +arguments + obj + f(:,:,:,:) + options.Staggering(1,:) char {mustBeMember(options.Staggering, {'none', 'forward', 'backward'})} = 'none' +end + +% Check input dimensions. +if obj.dimensions ~= size(f, 4) + kwave.toolbox.Logger.error('FourierCollocation:incorrectSize', ['Input must be vector field with ' num2str(obj.dimensions) ' components.']); +end + + +% Assign pseudonym for k-space derivative and shift operator. +switch options.Staggering + case 'none' + dfxdx = obj.ddxNoShift; + dfxdy = obj.ddyNoShift; + dfxdz = obj.ddzNoShift; + dfydx = obj.ddxNoShift; + dfydy = obj.ddyNoShift; + dfydz = obj.ddzNoShift; + dfzdx = obj.ddxNoShift; + dfzdy = obj.ddyNoShift; + dfzdz = obj.ddzNoShift; + case 'forward' + dfxdx = obj.ddxShiftPos; + dfxdy = obj.ddyShiftNeg; + dfxdz = obj.ddzShiftNeg; + dfydx = obj.ddxShiftNeg; + dfydy = obj.ddyShiftPos; + dfydz = obj.ddzShiftNeg; + dfzdx = obj.ddxShiftNeg; + dfzdy = obj.ddyShiftNeg; + dfzdz = obj.ddzShiftPos; + case 'backward' + dfxdx = obj.ddxShiftNeg; + dfxdy = obj.ddyShiftPos; + dfxdz = obj.ddzShiftPos; + dfydx = obj.ddxShiftPos; + dfydy = obj.ddyShiftNeg; + dfydz = obj.ddzShiftPos; + dfzdx = obj.ddxShiftPos; + dfzdy = obj.ddyShiftPos; + dfzdz = obj.ddzShiftNeg; +end + +% Preallocate output matrix (tensor field). +df = zeros([obj.kgridPadded.gridSize, obj.dimensions, obj.dimensions]); + +% Scalar or no k-space correction, so use 1D FFTs. +if isempty(obj.kappa) || isscalar(obj.kappa) + for dimInd = 1:obj.dimensions + switch dimInd + case 1 + df(:, :, :, 1, 1) = ifft(dfxdx .* fft(f(:, :, :, 1), [], 1), [], 1, 'symmetric'); + case 2 + df(:, :, :, 2, 1) = ifft(dfydx .* fft(f(:, :, :, 2), [], 1), [], 1, 'symmetric'); + df(:, :, :, 1, 2) = ifft(dfxdy .* fft(f(:, :, :, 1), [], 2), [], 2, 'symmetric'); + df(:, :, :, 2, 2) = ifft(dfydy .* fft(f(:, :, :, 2), [], 2), [], 2, 'symmetric'); + case 3 + df(:, :, :, 3, 1) = ifft(dfzdx .* fft(f(:, :, :, 3), [], 1), [], 1, 'symmetric'); + df(:, :, :, 3, 2) = ifft(dfzdy .* fft(f(:, :, :, 3), [], 2), [], 2, 'symmetric'); + df(:, :, :, 1, 3) = ifft(dfxdz .* fft(f(:, :, :, 1), [], 3), [], 3, 'symmetric'); + df(:, :, :, 2, 3) = ifft(dfydz .* fft(f(:, :, :, 2), [], 3), [], 3, 'symmetric'); + df(:, :, :, 3, 3) = ifft(dfzdz .* fft(f(:, :, :, 3), [], 3), [], 3, 'symmetric'); + end + end + + if isscalar(obj.kappa) + df = df .* obj.kappa; + end + +% ND k-space correction, so use ND FFTs. +else + for dimInd = 1:obj.dimensions + switch dimInd + case 1 + df(:, :, :, 1, 1) = ifftn(dfxdx .* obj.kappa .* fftn(f(:, :, :, 1)), 'symmetric'); + case 2 + df(:, :, :, 2, 1) = ifftn(dfydx .* obj.kappa .* fftn(f(:, :, :, 2)), 'symmetric'); + df(:, :, :, 1, 2) = ifftn(dfxdy .* obj.kappa .* fftn(f(:, :, :, 1)), 'symmetric'); + df(:, :, :, 2, 2) = ifftn(dfydy .* obj.kappa .* fftn(f(:, :, :, 2)), 'symmetric'); + case 3 + df(:, :, :, 3, 1) = ifftn(dfzdx .* obj.kappa .* fftn(f(:, :, :, 3)), 'symmetric'); + df(:, :, :, 3, 2) = ifftn(dfzdy .* obj.kappa .* fftn(f(:, :, :, 3)), 'symmetric'); + df(:, :, :, 1, 3) = ifftn(dfxdz .* obj.kappa .* fftn(f(:, :, :, 1)), 'symmetric'); + df(:, :, :, 2, 2) = ifftn(dfydz .* obj.kappa .* fftn(f(:, :, :, 2)), 'symmetric'); + df(:, :, :, 3, 3) = ifftn(dfzdz .* obj.kappa .* fftn(f(:, :, :, 3)), 'symmetric'); + end + end +end