Small additions to the calculator capabilities

TZPGcalc.py

 # -*- coding: utf-8 -*-
 """ TZPG simple calculator for SCS.
 
-    Interactive widget to calculate beam sizes and position at the sample and detector planes for the SCS instrument.
+    Interactive widget to calculate beam sizes and position at the sample and
+    detector planes for the SCS instrument.
 
     Copyright (2019) SCS Team.
 """
@@ -26,20 +27,23 @@ d = 3.3 - Z0 # distance between HFM and TZPG
 f1 = 7.3 # HFM focus 2 m behind second interaction point
 F = F*(d-f1)/(d-f1-F)
 
+# number of membrane to show
+SampleN = 7
+
 class TZPGcalc():
     def __init__(self):
         self.initFig()
         self.initWidgets()
         self.UpdateFig()
         display(self.control)
-        
+
     def initFig(self):
         """ Creates a figure for the sample plane and detector plane images with all necessary drawings.
         """
-        
+
         plt.close('TZPGcalc')
         fig, (self.ax_sam, self.ax_det) = plt.subplots(1, 2, num='TZPGcalc', figsize=(6,3))
-        
+
         # display scale
         self.scale = 1e3 # displayed distances in [mm]
 
@@ -90,7 +94,7 @@ class TZPGcalc():
                 'F1G1': self.ax_det.add_patch(Rectangle((0, 0), 1, 1, facecolor=c_gb, alpha=0.7, lw=None)),
                 'F1G-1': self.ax_det.add_patch(Rectangle((0, 0), 1, 1, facecolor=c_gb, alpha=0.7, lw=None))
                }
-        
+
         self.detLines = {'module': self.ax_det.add_patch(Rectangle((0, 0), 1, 1, fill=False, facecolor='k')),
                 'Vfilter': self.ax_det.add_patch(Rectangle((0, 0), 1, 1, facecolor="blue", alpha=0.4)),
                 'Hfilter': self.ax_det.add_patch(Rectangle((0, 0), 1, 1, facecolor="blue", alpha=0.4)),
@@ -100,31 +104,31 @@ class TZPGcalc():
         # 5x5 membranes
         self.sampleLines = {}
         self.etchLines = {}
-        for k in range(25):
+        for k in range(SampleN*SampleN):
             self.sampleLines[k] = self.ax_sam.add_patch(Rectangle((0, 0), 1, 1, fill=False, facecolor='k'))
             self.etchLines[k] = self.ax_sam.add_patch(Rectangle((0, 0), 1, 1, fill=False, facecolor='k', alpha=0.4, ls='--'))
 
     def RectUpdate(self, rect, xLeft, yBottom, xRight, yTop):
         """ Updates the position and size of the given Rectangle.
-        
+
             rect: Rectangle to update
             xLeft: x position of the left corner
             yBottom: y position of the bottom corner
             xRight: x position of the right corner
             yTop: y position of the top corner
         """
-        
+
         xw = np.abs(xLeft - xRight)
         yw = np.abs(yTop - yBottom)
-        
-    
+
+
         rect.set_xy((self.scale*xLeft, self.scale*yBottom))
         rect.set_height(self.scale*yw)
         rect.set_width(self.scale*xw)
-    
+
     def UpdateBeams(self, Beams, Z, conf):
         """ Update the position and size of the beams.
-        
+
             Beams: dictionary of f'F{f}G{g}' Rectangles for f = 0 and 1 zone plate order and
                 g = +1, 0 and -1 grating order
             Z: distance Z between the zone plate and the current imaging plane
@@ -148,7 +152,7 @@ class TZPGcalc():
         YdfL = (o*(Z - F)/F + offaxis)
         Ydf = ((o + wV/2)*(Z - F)/F + offaxis)
         YdfH = ((o + wV)*(Z - F)/F + offaxis)
-        
+
         if Z < F: # before zone plate focus, low and high beam edges are swapped
             Xdg1L, Xdg1H = (Xdg1H, Xdg1L)
             YdfL, YdfH = (YdfH, YdfL)
@@ -169,7 +173,7 @@ class TZPGcalc():
 
     def DetectorUpdate(self, Xoff, Yoff):
         """ Draw DSSC detector module with filter mask.
-        
+
             Xoff: x offset
             Yoff: y offset
         """
@@ -191,13 +195,13 @@ class TZPGcalc():
                    -filterW/2 + Xoff, -filterL/2 + Yoff,  filterW/2 + Xoff, filterL/2 + Yoff)
         self.RectUpdate(self.detLines['Hfilter'],
                    -filterL/2 + Xoff, -filterW/2 + Yoff,  filterL/2 + Xoff, filterW/2 + Yoff)
-        
+
         # moving rotated rectangles is a pain in matplotlib
         self.detLines['diamond'].set_xy((self.scale*Xoff, self.scale*(Yoff - diamondW/2*np.sqrt(2))))
 
     def SampleUpdate(self, w, p, Xoff, Yoff, thickness=0.525):
         """ Draw the sample.
-        
+
             w: membrane width
             p: membrane pitch
             Xoff: sample x offset
@@ -208,22 +212,23 @@ class TZPGcalc():
         wp = w +2*thickness/np.tan(np.deg2rad(54.74))
 
         j = 0
-        for k in range(-2, 3):
-            for l in range(-2, 3):
+        for k in range(-(SampleN-1)//2, (SampleN-1)//2+1):
+            for l in range(-(SampleN-1)//2, (SampleN-1)//2+1):
                 self.RectUpdate(self.sampleLines[j], k*p - w/2 + Xoff, l*p - w/2 - Yoff, k*p + w/2 + Xoff, l*p + w/2 - Yoff)
                 self.RectUpdate(self.etchLines[j], k*p - wp/2 + Xoff, l*p - wp/2 - Yoff, k*p + wp/2 + Xoff, l*p + wp/2 - Yoff)
                 j+=1
 
     def UpdateFig(self):
         """ Update the figure with the current slider values.
-        
+
         """
-        
+
         # we calculate the optics for the central wavelength
         nrjL, nrjH = self.nrj_slider.value # [eV]
         wlL = 1240/nrjL*1e-9
         wlH = 1240/nrjH*1e-9
-        wl = 0.5*(wlL + wlH)
+        nrjD = self.design_nrj_slider.value # [eV]
+        wl = 1240/nrjD*1e-9
 
         theta_grating = self.grating_slider.value*1e-3 # [rad]
         sampleZ = self.samz_slider.value*1e-3 # [m]
@@ -281,7 +286,14 @@ class TZPGcalc():
         self.nrj_slider = widgets.FloatRangeSlider(
             value=[840., 880.],
             min=450.,
-            max=3000.0,
+            max=3200.0,
+            step=1,
+            readout_format='.2f',
+        )
+        self.design_nrj_slider = widgets.FloatSlider(
+            value=860.,
+            min=450.,
+            max=3200.0,
             step=1,
             readout_format='.2f',
         )
@@ -316,6 +328,7 @@ class TZPGcalc():
         self.dr_label = widgets.Label(value='dr')
         self.d_label = widgets.Label(value='dr')
         TZPGTab = VBox(children=[HBox([widgets.Label(value='Energy (eV):'), self.nrj_slider]),
+                                 HBox([widgets.Label(value='Design Energy (eV):'), self.design_nrj_slider]),
                                  HBox([widgets.Label(value=r'Grating $\theta$ (mrad):'), self.grating_slider]),
                                  self.d_label, self.dr_label,
                                  HBox([widgets.Label(value='TZPG horiz. width (mm):'), self.TZPGwH_slider]),
@@ -413,4 +426,4 @@ class TZPGcalc():
         tab3.set_title(0, 'detector')
         tab3.selected_index = None
 
-        self.control = VBox(children=[tab1, tab2, tab3, self.button])
+        self.control = VBox(children=[tab1, tab2, tab3, self.button])
\ No newline at end of file