Update

diffractsim/polychromatic_simulator.py

@@ -173,17 +173,19 @@ def get_colors_at_image_plane(self, pupil, M, zi, z0, scale_factor = 1):
         
         """
 
+
+
+        for j in range(len(self.optical_elements)):
+            self.E = self.E * self.optical_elements[j].get_transmittance(self.xx, self.yy, 0)
+
+
         # if the magnification is negative, the image is inverted
         if M < 0:
             self.E = bd.flip(self.E)
         M_abs = bd.abs(M)
 
         self.E = self.E/M_abs
 
-
-        for j in range(len(self.optical_elements)):
-            self.E = self.E * self.optical_elements[j].get_transmittance(self.xx, self.yy, 0)
-
         fft_c = bd.fft.fft2(self.E)
         c = bd.fft.fftshift(fft_c)