1 files changed, 10 insertions, 5 deletions
diff --git a/fitness.m b/fitness.m
index 54060a7..2cb2278 100644
--- a/fitness.m
+++ b/fitness.m
@@ -1,12 +1,17 @@
-function [Energy, Waist, Penalty] = fitness( q_0, q_final, x_final, optics_positions, optics_focal_length, lambda )
+function [Energy, Waist_backward, Waist_forward, Penalty] = fitness( q_0, q_final, x_final, optics_positions, optics_focal_length, lambda )
 %FITNESS Summary of this function goes here
 %   Detailed explanation goes here
     x0 = 0;
-    q_f_trial = gbeam_propagation(x_final,q_0,x0,optics_placer(optics_positions, optics_focal_length));
-    [Waist, Radius] = q2wr(q_f_trial, lambda);
+    x_forward=linspace(0,x_final,15);
+    q_forward = gbeam_propagation(x_forward,q_0,x0,optics_placer(optics_positions, optics_focal_length));
+    [Waist_forward, Radius_forward] = q2wr(q_forward, lambda);
         
-    Energy = abs(q_final-q_f_trial);
-    
+    x_backward=(x_forward);
+    q_backward = gbeam_propagation(x_backward,q_final,x_final,optics_placer(optics_positions, optics_focal_length));
+    [Waist_backward, Radius_backward] = q2wr(q_backward, lambda);
+        
+    Energy = sum((Waist_backward - Waist_forward).^2)/(mean(Waist_backward))^2;
+           
     % penalty calculation
     % do not put lenses too close to each other and end positions
     lens_size=0.03;