1 files changed, 56 insertions, 8 deletions

diff --git a/solution_visualization.m b/solution_visualization.m
index 2cf8ce7..2b4158c 100644
--- a/solution_visualization.m
+++ b/solution_visualization.m
@@ -1,11 +1,19 @@
-function [waste_at_the_end, radius_at_the_end] = solution_visualization(q0,x0, qf, xf, optics, lambda)
+function [waste_at_the_end, radius_at_the_end, waist_at_lens_position] = solution_visualization(q0,x0, qf, xf, optics, lambda, lens_width, display_prop)
 %Propagates beam with given input parameters
 %   Forward propagation and backward propagation are taken in order to
 %   visualize reasonable solutions.
 
+%Array of lens positions
+n_lens = size(optics,2);
+
+for i = 1:n_lens
+    lens_position(i) = optics{i}.x;
+end
+
 x=linspace(x0,xf,1000); % we will calculate beam profile between x0 and xf
+
 %Forward propagation
-q_forward=gbeam_propagation(x,q0,x0,optics);
+q_forward =gbeam_propagation(x,q0,x0,optics);
 [w_forward,r_forward]=q2wr(q_forward, lambda);
 
 %Backward propagation
@@ -13,13 +21,53 @@ q_backward=gbeam_propagation(x,qf,xf,optics);
 [w_backward,r_backward]=q2wr(q_backward, lambda);
 
 %Plot beam profile
-plot ( ...
-	x,w_forward, '-r',  ...
-	x,-w_forward, '-r', ...
-	x, w_backward, '-.b', ...
-	x, -w_backward, '-.b')
-legend({'forward propagation', '', 'backward propagation', ''})      
+subplot(2,1,1); plot ( ...
+    x,w_forward, '-r',  ...
+    x,-w_forward, '-r', ...
+    x, w_backward, '-.b', ...
+    x, -w_backward, '-.b')
+legend({'forward propagation', '', 'backward propagation', ''})
+
+%Find q and waist at lens positions
+waist_at_lens_position = zeros(1,n_lens);
+q_lens = zeros(1,n_lens);
+
+for i = 1:n_lens
+    q_lens(i) = gbeam_propagation(lens_position(i),q0,x0,optics);
+end
+
+
+for i = 1:n_lens
+    waist_at_lens_position(i) = q2wr(q_lens(i), lambda);
+end
+
+
+%Plot lenses
+color = ['m' 'g' 'c'];
+
+if display_prop(1) == 1
+    for i = 1:n_lens
+        half_width = lens_width/2;
+        corrected_lens_position = lens_position(i)-half_width;
+        corrected_waist = waist_at_lens_position(i)*2;
+        rectangle('Position', [corrected_lens_position,-waist_at_lens_position(i),lens_width,corrected_waist], 'EdgeColor', color(i), 'LineWidth',1);
+    end
+end
+
+if display_prop(2) == 1
+    for i = 1:n_lens
+        x1 = optics{i}.x;
+        y1 = waist_at_lens_position(i);
+        x2 = x1;
+        y2 = -waist_at_lens_position(i);
+        
+        line([x1;x2], [y1;y2], 'LineWidth', 3, 'Color', color(i));
+    end
+end
+
 
 [waste_at_the_end,radius_at_the_end]  = q2wr(q_forward(end), lambda);
 
 
+
+