Extended sum support for arrays of beams

ilent2 · Mar 26, 2020 · 4d8ebc3 · 4d8ebc3
1 parent 9b05cde
commit 4d8ebc3
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 14 deletions.
diff --git a/+ott/Bsc.m b/+ott/Bsc.m
@@ -21,7 +21,7 @@
 %
 % Methods
 %   - append      --  Joins two beam objects together
-%   - mergeBeams  --  Merge the BSCs for the beams contained in this object
+%   - sum         --  Merge the BSCs for the beams contained in this object
 %   - translateZ  --  Translates the beam along the z axis
 %   - translateXyz -- Translation to xyz using rotations and z translations
 %   - translateRtp -- Translation to rtp using rotations and z translations
@@ -352,13 +352,6 @@
         beam.b = [beam.b, other.b];
       end
     end
-
-    function beam = mergeBeams(beam)
-      % Merge the BSCs for the beams contained in this object
-
-      beam.a = sum(beam.a, 2);
-      beam.b = sum(beam.b, 2);
-    end
 
     function [E, H, data] = farfield(beam, theta, phi, varargin)
       %FARFIELD finds far field at locations theta, phi.
@@ -1775,17 +1768,52 @@ function visualiseFarfieldSphere(beam, varargin)
       beam.b = beam.b - beam2.b;
     end
 
-    function beam = sum(beam1)
+    function beam = sum(beamin, dim)
       % Sum beam coefficients
       %
       % Usage
       %   beam = sum(beam)
       %
       %   beam = beam.sum()
+      %
+      %   beam = sum([beam1, beam2, ...], dim) sums the given beams,
+      %   similar to Matlab's ``sum`` builtin.  ``dim`` is the dimension
+      %   to sum over (optional).
 
-      beam = beam1;
-      beam.a = sum(beam.a, 2);
-      beam.b = sum(beam.b, 2);
+      if numel(beamin) > 1
+        % beam is an array
+
+        % Handle default value for dimension
+        if nargin < 2
+          if isvector(beamin)
+            dim = find(size(beamin) > 1, 1);
+          elseif ismatrix(beamin) == 2
+            dim = 2;
+          else
+            dim = find(size(beamin) > 1, 1);
+          end
+        end
+
+        % Select the first row in our dimension
+        subs = [repmat({':'},1,dim-1), 1, ...
+          repmat({':'},1,ndims(beamin)-dim)];
+        S = struct('type', '()', 'subs', {subs});
+        beam = subsref(beamin, S);
+
+        % Add each beam
+        for ii = 2:size(beamin, dim)
+        subs = [repmat({':'},1,dim-1), ii, ...
+          repmat({':'},1,ndims(beamin)-dim)];
+          S = struct('type', '()', 'subs', {subs});
+          beam = beam + subsref(beamin, S);
+        end
+
+      else
+        % Beam union
+        beam = beamin;
+        beam.a = sum(beam.a, 2);
+        beam.b = sum(beam.b, 2);
+      end
     end
 
     function beam = clearDz(beam)

diff --git a/tests/testBsc.m b/tests/testBsc.m
@@ -137,7 +137,15 @@ function testSum(testCase)
   beam3 = beam1 + beam2;
 
   beamU = beam1.append(beam2);
-  testCase.verifyEqual(beamU.sum(), beam3);
-  testCase.verifyEqual(sum(beamU), beam3);
+  testCase.verifyEqual(beamU.sum(), beam3, ...
+    'beamU.sum() incorrect');
+  testCase.verifyEqual(sum(beamU), beam3, ...
+    'sum(beamU) incorrect');
+
+  % Test array sum
+  beamarr = [beam1, beam2];
+  testCase.verifyEqual(sum(beamarr), beam3, ...
+    'sum([beam1, beam2]) incorrect');
 
 end
+