diff --git a/+kwave/+tests/+unit/TestFourierCollocation.m b/+kwave/+tests/+unit/TestFourierCollocation.m index 6a9ba65d..e494f115 100644 --- a/+kwave/+tests/+unit/TestFourierCollocation.m +++ b/+kwave/+tests/+unit/TestFourierCollocation.m @@ -87,6 +87,37 @@ function testDivergence(testCase) end + % Test the gradient of a vector function. + function testDivergenceTensorSplit(testCase) + import matlab.unittest.constraints.IsEqualTo + + % No staggering. + [f, testCase.referenceSolution] = testCase.getPeriodicGradTensorFunction; + testCase.actualSolution = testCase.solver.divergenceTensorSplit(f); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Forward staggering. + [f, testCase.referenceSolution] = testCase.getPeriodicGradTensorFunction("forward"); + testCase.actualSolution = testCase.solver.divergenceTensorSplit(f, Staggering="forward"); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Backward staggering. + [f, testCase.referenceSolution] = testCase.getPeriodicGradTensorFunction("backward"); + testCase.actualSolution = testCase.solver.divergenceTensorSplit(f, Staggering="backward"); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Scalar kappa. + testCase.solver.kappa = 2; + [f, testCase.referenceSolution] = testCase.getPeriodicGradTensorFunction; + testCase.referenceSolution = testCase.referenceSolution .* testCase.solver.kappa; + testCase.actualSolution = testCase.solver.divergenceTensorSplit(f); + testCase.verifyThat(testCase.actualSolution, IsEqualTo(testCase.referenceSolution, "Within", testCase.tol)); + + % Test incorrect size gives exception. + f = rand(2, 2, 2, 7); + testCase.verifyError(@() testCase.solver.divergenceTensorSplit(f), 'FourierCollocation:incorrectSize'); + end + % Test the curl function. function testCurl(testCase) import matlab.unittest.constraints.IsEqualTo @@ -303,17 +334,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 +358,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 +368,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 +394,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,34 +415,141 @@ 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 + + % Define a periodic tensor function and its analytic gradient on + % the grid specified by obj.kgridPadded, returning the gradients + % in each axis. The function is normalized so the maximum of the + % gradient in each axis is approximately 1. The gradient tensor + % can also be returned on a staggered grid. + function [F, gradF] = getPeriodicGradTensorFunction(obj, staggering) + + arguments + obj + staggering(1,:) char {mustBeMember(staggering, {'none', 'forward', 'backward'})} = 'none' + end + + switch staggering + case 'none' + xxSgx = obj.kgridPadded.xVec; + xySgy = obj.kgridPadded.yVec; + xzSgz = obj.kgridPadded.zVec; + yxSgx = obj.kgridPadded.xVec; + yySgy = obj.kgridPadded.yVec; + yzSgz = obj.kgridPadded.zVec; + zxSgx = obj.kgridPadded.xVec; + zySgy = obj.kgridPadded.yVec; + zzSgz = obj.kgridPadded.zVec; + case 'forward' + xxSgx = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + xySgy = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + xzSgz = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + yxSgx = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + yySgy = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + yzSgz = obj.kgridPadded.zVec - obj.kgridPadded.dz/2; + zxSgx = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + zySgy = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + zzSgz = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + + case 'backward' + xxSgx = obj.kgridPadded.xVec - obj.kgridPadded.dx/2; + xySgy = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + xzSgz = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + yxSgx = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + yySgy = obj.kgridPadded.yVec - obj.kgridPadded.dy/2; + yzSgz = obj.kgridPadded.zVec + obj.kgridPadded.dz/2; + zxSgx = obj.kgridPadded.xVec + obj.kgridPadded.dx/2; + zySgy = obj.kgridPadded.yVec + obj.kgridPadded.dy/2; + zzSgz = 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 .* xxSgx); + + case 2 + kx = (2*pi ./ obj.kgridPadded.xSize); + ky = (2*pi ./ obj.kgridPadded.ySize); + + [X, Y] = ndgrid(obj.kgridPadded.xVec, obj.kgridPadded.yVec); + [xxsgx, xysgy] = ndgrid(xxSgx, xySgy); + [yxsgx, yysgy] = ndgrid(yxSgx, yySgy); + + Fxx = sin(kx .* X) .* sin(ky .* Y) ./ kx; + Fyy = sin(ky .* Y) .* sin(kx .* X) ./ ky; + Fxy = sin(kx .* X) .* sin(ky .* Y) ./ (kx .* ky); + + gradFxx_x = cos(kx .* xxsgx) .* sin(ky .* Y); + gradFxy_x = sin(ky .* Y) .* cos(kx .* yxsgx) .* (1 ./ ky); + gradFxy_y = sin(kx .* X) .* cos(ky .* xysgy) .* (1 ./ kx); + gradFyy_y = cos(ky .* yysgy) .* sin(kx .* X); + + F = cat(4, Fxx, Fyy, Fxy); + gradF = cat(5, cat(4, gradFxx_x, gradFxy_x), cat(4, gradFxy_y, gradFyy_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); + [xxsgx, xysgy, xzsgz] = ndgrid(xxSgx, xySgy, xzSgz); + [yxsgx, yysgy, yzsgz] = ndgrid(yxSgx, yySgy, yzSgz); + [zxsgx, zysgy, zzsgz] = ndgrid(zxSgx, zySgy, zzSgz); + + Fxx = sin(kx .* X) .* sin(ky .* Y) .* sin(kz .* Z) ./ kx; + Fyy = sin(ky .* Y) .* sin(kx .* X) .* sin(kz .* Z) ./ ky; + Fzz = sin(kz .* Z) .* sin(kx .* X) .* sin(ky .* Y) ./ kz; + Fxy = sin(kx .* X) .* sin(ky .* Y) .* sin(kz .* Z) ./ (kx .* ky); + Fxz = sin(ky .* Y) .* sin(kx .* X) .* sin(kz .* Z) ./ (kx .* kz); + Fyz = sin(kz .* Z) .* sin(kx .* X) .* sin(ky .* Y) ./ (ky .* kz); + + gradFxx_x = cos(kx .* xxsgx) .* sin(ky .* Y) .* sin(kz .* Z); + gradFxy_y = sin(kx .* X) .* cos(ky .* xysgy) .* sin(kz .* Z) .* (1 ./ kx); + gradFxz_z = sin(kx .* X) .* sin(ky .* Y) .* cos(kz .* xzsgz) .* (1 ./ kx); + + gradFyx_x = sin(ky .* Y) .* cos(kx .* yxsgx) .* sin(kz .* Z) .* (1 ./ ky); + gradFyy_y = cos(ky .* yysgy) .* sin(kx .* X) .* sin(kz .* Z); + gradFyz_z = sin(ky .* Y) .* sin(kx .* X) .* cos(kz .* yzsgz) .* (1 ./ ky); + + gradFzx_x = sin(kz .* Z) .* cos(kx .* zxsgx) .* sin(ky .* Y) .* (1 ./ kz); + gradFzy_y = sin(kz .* Z) .* sin(kx .* X) .* cos(ky .* zysgy) .* (1 ./ kz); + gradFzz_z = cos(kz .* zzsgz) .* sin(kx .* X) .* sin(ky .* Y); + + F = cat(4, Fxx, Fyy, Fzz, Fxy, Fxz, Fyz); + gradF = cat(5, cat(4, gradFxx_x, gradFyx_x, gradFzx_x), cat(4, gradFxy_y, gradFyy_y, gradFzy_y), cat(4, gradFxz_z, gradFyz_z, gradFzz_z)); + end + end + % Define a periodic scalar function and its analytic laplacian on % the grid specified by obj.kgridPadded. The function is normalised % so the maximum of the gradient is approximately 1. The Laplacian diff --git a/+kwave/+toolbox/@FourierCollocation/divergenceTensorSplit.m b/+kwave/+toolbox/@FourierCollocation/divergenceTensorSplit.m new file mode 100644 index 00000000..4d54ed46 --- /dev/null +++ b/+kwave/+toolbox/@FourierCollocation/divergenceTensorSplit.m @@ -0,0 +1,174 @@ +%% divergenceTensorSplit +% *Class:* kwave.toolbox.FourierCollocation +% *Package:* kwave.toolbox +% +% Calculate the divergence of a tensor field. +% +%% Syntax +% ds = divergenceTensorSplit(obj, s) +% ds = divergenceTensorSplit(obj, s, Staggering='forward') +% +%% Description +% Calculates the divergence of a symmetric tensor in 1D, 2D, or 3D using a +% Fourier collocation spectral method. The symmetric tensor must be +% linearised to a vector of dimensions % [D * (D + 1) / 2 , 1] where the +% components are [sxx] in 1D, [sxx, syy, sxy]' in 2D and +% [sxx, syy, szz, sxy, sxz, syz]' in 3D. +% +% The partial derivatives of the symmetric tensor which are non-zero are +% ds = [dsxxdx, dsxydy, dsxzdz; +% dsxydx, dsyydy, dsyzdz; +% dsxzdx, dsyzdy, dszzdz] +% +% The tensor components of the divergence are stacked in the 4th and 5th +% dimension of the output. For example, if calling divergenceTensorSplit +% 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. +% +% If obj.kappa is defined, a k-space correction is applied as part of the +% divergence calculation. If kappa is a scalar (single frequency correction) +% or empty, the divergence components are calculated using 1D FFTs. If kappa +% is a matrix, the divergence components are calculated using ND FFTs, and +% kappa is applied in the Fourier domain. +% +% The divergence 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 +% * |s| - (numeric) Tensor field to compute the split divergence 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 +% * |ds| - (numeric) split divergence of the tensor field s. + +% 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 ds = divergenceTensorSplit(obj, s, options) + +arguments + obj + s(:,:,:,:) + options.Staggering(1,:) char {mustBeMember(options.Staggering, {'none', 'forward', 'backward'})} = 'none' +end + +% Check input dimensions. +expectedSize = obj.dimensions * (obj.dimensions + 1) / 2; +if expectedSize ~= size(s, 4) + kwave.toolbox.Logger.error('FourierCollocation:incorrectSize', ['Input must be vector field with ' num2str(expectedSize) ' components.']); +end + +% Assign pseudonym for k-space derivative and shift operator. +switch options.Staggering + case 'none' + dsxxdx = obj.ddxNoShift; + dsxydy = obj.ddyNoShift; + dsxzdz = obj.ddzNoShift; + dsxydx = obj.ddxNoShift; + dsyydy = obj.ddyNoShift; + dsyzdz = obj.ddzNoShift; + dsxzdx = obj.ddxNoShift; + dsyzdy = obj.ddyNoShift; + dszzdz = obj.ddzNoShift; + case 'forward' + dsxxdx = obj.ddxShiftPos; + dsxydy = obj.ddyShiftNeg; + dsxzdz = obj.ddzShiftNeg; + dsxydx = obj.ddxShiftNeg; + dsyydy = obj.ddyShiftPos; + dsyzdz = obj.ddzShiftNeg; + dsxzdx = obj.ddxShiftNeg; + dsyzdy = obj.ddyShiftNeg; + dszzdz = obj.ddzShiftPos; + case 'backward' + dsxxdx = obj.ddxShiftNeg; + dsxydy = obj.ddyShiftPos; + dsxzdz = obj.ddzShiftPos; + dsxydx = obj.ddxShiftPos; + dsyydy = obj.ddyShiftNeg; + dsyzdz = obj.ddzShiftPos; + dsxzdx = obj.ddxShiftPos; + dsyzdy = obj.ddyShiftPos; + dszzdz = obj.ddzShiftNeg; +end + +% Preallocate output matrix (vector field). +ds = 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 + ds(:, :, :, 1, 1) = ifft(dsxxdx .* fft(s(:, :, :, 1), [], 1), [], 1, 'symmetric'); + case 2 + ds(:, :, :, 2, 2) = ifft(dsyydy .* fft(s(:, :, :, 2), [], 2), [], 2, 'symmetric'); + if obj.dimensions == 2 + ds(:, :, :, 2, 1) = ifft(dsxydx .* fft(s(:, :, :, 3), [], 1), [], 1, 'symmetric'); + ds(:, :, :, 1, 2) = ifft(dsxydy .* fft(s(:, :, :, 3), [], 2), [], 2, 'symmetric'); + elseif obj.dimensions == 3 + ds(:, :, :, 2, 1) = ifft(dsxydx .* fft(s(:, :, :, 4), [], 1), [], 1, 'symmetric'); + ds(:, :, :, 1, 2) = ifft(dsxydy .* fft(s(:, :, :, 4), [], 2), [], 2, 'symmetric'); + end + case 3 + ds(:, :, :, 3, 3) = ifft(dszzdz .* fft(s(:, :, :, 3), [], 3), [], 3, 'symmetric'); + ds(:, :, :, 3, 1) = ifft(dsxzdx .* fft(s(:, :, :, 5), [], 1), [], 1, 'symmetric'); + ds(:, :, :, 1, 3) = ifft(dsxzdz .* fft(s(:, :, :, 5), [], 3), [], 3, 'symmetric'); + ds(:, :, :, 3, 2) = ifft(dsyzdy .* fft(s(:, :, :, 6), [], 2), [], 2, 'symmetric'); + ds(:, :, :, 2, 3) = ifft(dsyzdz .* fft(s(:, :, :, 6), [], 3), [], 3, 'symmetric'); + end + end + + if isscalar(obj.kappa) + ds = ds .* obj.kappa; + end + +% ND k-space correction, so use ND FFTs. +else + for dimInd = 1:obj.dimensions + switch dimInd + case 1 + ds(:, :, :, 1, 1) = ifftn(dsxxdx .* obj.kappa .* fftn(s(:, :, :, 1)), 'symmetric'); + case 2 + ds(:, :, :, 2, 2) = ifftn(dsyydy .* obj.kappa .* fftn(s(:, :, :, 2)), 'symmetric'); + if obj.dimensions == 2 + ds(:, :, :, 2, 1) = ifftn(dsxydx .* obj.kappa .* fftn(s(:, :, :, 3)), 'symmetric'); + ds(:, :, :, 1, 2) = ifftn(dsxydy .* obj.kappa .* fftn(s(:, :, :, 3)), 'symmetric'); + elseif obj.dimensions == 3 + ds(:, :, :, 2, 1) = ifftn(dsxydx .* obj.kappa .* fftn(s(:, :, :, 4)), 'symmetric'); + ds(:, :, :, 1, 2) = ifftn(dsxydy .* obj.kappa .* fftn(s(:, :, :, 4)), 'symmetric'); + end + case 3 + ds(:, :, :, 3, 3) = ifftn(dszzdz .* obj.kappa .* fftn(s(:, :, :, 3)), 'symmetric'); + ds(:, :, :, 3, 1) = ifftn(dsxzdx .* obj.kappa .* fftn(s(:, :, :, 5)), 'symmetric'); + ds(:, :, :, 1, 3) = ifftn(dsxzdz .* obj.kappa .* fftn(s(:, :, :, 5)), 'symmetric'); + ds(:, :, :, 3, 2) = ifftn(dsyzdy .* obj.kappa .* fftn(s(:, :, :, 6)), 'symmetric'); + ds(:, :, :, 2, 2) = ifftn(dsyzdz .* obj.kappa .* fftn(s(:, :, :, 6)), 'symmetric'); + end + end +end