Merge pull request #112 from samoylv/screen

Screen

s2e/prop/exfel_spb_day1_beamline.py

@@ -5,6 +5,7 @@
 from opd import defineOPD
 import os
 
+from prop.propagate_s2e import MIRROR_DATA_DIR as mirror_data_dir
 
 def get_beamline():
     """ Setup and return the WPG.Beamline object representing the SPB/SFX nanofocus beamline (KB mirrors).

s2e/prop/propagate_s2e.py

@@ -1,7 +1,7 @@
 """ Module for interface functions to wpg library. """
 
 from glob import glob
-from wpg import Wavefront
+from wpg import Wavefront, Beamline
 import errno
 import h5py
 import multiprocessing
@@ -135,9 +135,95 @@ def propagate(in_fname, out_fname, get_beamline):
     #Resizing: decreasing Range of Horizontal and Vertical Position:
     wpg.srwlib.srwl.ResizeElecField(wf._srwl_wf, 'c', [0, 0.5, 1, 0.5,  1]);
 
+    add_custom_data(wf, bl0)
+
+    print('Saving propagated wavefront to ' + out_fname)
+    mkdir_p(os.path.dirname(out_fname))
+    wf.store_hdf5(out_fname)
+
+    print('Saving history.')
+    add_history(out_fname, in_fname)
+
+    print('ALL DONE.')
+    print('#'*80)
+
+
+
+def stepwise(in_fname, get_beamline):
+    """
+    Propagate wavefront stepwise, dumping the wavefront at every step.
+
+    :param in_file: input wavefront file
+    :param get_beamline: function to build beamline
+    """
+
+    print('#'*80)
+    print("Setup initial wavefront.")
+    wf=Wavefront()
+
+    # Load wavefront data.
+    print('Load ' + in_fname)
+    wf.load_hdf5(in_fname)
+
+    # Get beamline.
+    bl0 = get_beamline()
+
+    beamline = bl0.propagation_options
+    if len(beamline) > 1:
+        raise RuntimeError("Beamline configuration not supported.")
+    beamline = beamline[0]
+    elements = beamline['optical_elements']
+    options = beamline['propagation_parameters']
+    if len(elements) != len(options):
+        raise RuntimeError("Beamline configuration not supported.")
+
+    i = 0
+    for element, option in zip(elements, options):
+
+        print('\n')
+        print('#'*80)
+        print('Propagation step %d.' %(i))
+        print('Setting up incremental beamline.')
+        beamline_step = Beamline()
+        beamline_step.append(element, option) ### <== CHECKME
+
+        # Switch to frequency domain.
+        wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 'f')
+
+        # Save spectrum for later reference.
+        sz0 = get_intensity_on_axis(wf);
+        wf.custom_fields['/misc/spectrum0'] = sz0
+
+        # Propagate.
+        beamline_step.propagate(wf)
+
+        # Save spectrum after propagation for later reference.
+        sz1 = get_intensity_on_axis(wf);
+        wf.custom_fields['/misc/spectrum1'] = sz1
+
+        # Switch back to time domain.
+        wpg.srwlib.srwl.SetRepresElecField(wf._srwl_wf, 't')
+
+        incremental_filename = "%04d.h5" % (i)
+        print('Saving propagated wavefront to ' + incremental_filename)
+        mkdir_p(os.path.dirname(incremental_filename))
+        wf.store_hdf5(incremental_filename)
+
+        print('Done with propagation step %d.' %(i))
+        print('#'*80)
+
+        # Increment running index.
+        i += 1
+
+
+
+def add_custom_data(wf, bl0):
+    """ Utility to store additional info on the wavefront . """
+
     # Add more metadata.
     fwhm = calculate_fwhm(wf)
 
+
     wf.custom_fields['/misc/xFWHM'] = fwhm['fwhm_x']
     wf.custom_fields['/misc/yFWHM'] = fwhm['fwhm_y']
     wf.custom_fields['/params/beamline/printout'] = str(bl0)
@@ -149,15 +235,6 @@ def propagate(in_fname, out_fname, get_beamline):
     wf.custom_fields['/info/method_description'] = """WPG, WaveProperGator (http://github.com/samoylv/WPG)is an interactive simulation framework for coherent X-ray wavefront propagation.\nSRW, Synchrotron Radiation Workshop (http://github.com/ochubar/SRW),  is a physical optics computer code  for simulation of the radiation wavefront propagation through optical systems of beamlines as well as  detailed characteristics of Synchrotron Radiation (SR) generated by relativistic electrons in magnetic fields of arbitrary configuration."""
     wf.custom_fields['/info/package_version'] = '2014.1'
 
-    print('Saving propagated wavefront to ' + out_fname)
-    mkdir_p(os.path.dirname(out_fname))
-    wf.store_hdf5(out_fname)
-
-    print('Saving history.')
-    add_history(out_fname, in_fname)
-
-    print('ALL DONE.')
-    print('#'*80)
 
 def propagate_wrapper(params):
     """

s2e/prop/simple_beamline.py

@@ -1,7 +1,7 @@
 def get_beamline():
     """
     This function will called in propagation script.
-    
+
     :return: Beamline.
     """
 
@@ -21,4 +21,4 @@ def get_beamline():
     bl0.append(lens0,  Use_PP())
     bl0.append(drift1, Use_PP(semi_analytical_treatment=1, zoom=4.2,  sampling=0.5))
 
-    return bl0
+    return bl0

s2e/prop/unittest/BeamlinesTest.py

@@ -35,7 +35,7 @@
 sys.path.insert(0,'../../..')
 
 
-from s2e.prop import simple_beamline, exfel_spb_day1_beamline, exfel_spb_kb_beamline
+from s2e.prop import simple_beamline, exfel_spb_day1_beamline, exfel_spb_kb_beamline, beamline_with_screen
 from s2e.prop import propagate_s2e
 
 from wpg import Beamline, Wavefront
@@ -130,6 +130,19 @@ def testEXFELSPBKBBeamlinePropagation(self):
         # Check that output was generated.
         self.assertIn( output_file, os.listdir(".") )
 
+    def testStepwise(self):
+        """ Test stepwise propagation through a beamline."""
+
+        propagate_s2e.stepwise(self.__fel_source, simple_beamline.get_beamline )
+
+        # Should produce 3 files.
+        for f in range(3):
+            filename = "%04d.h5" % (f)
+            self.__files_to_remove.append(filename)
+            self.assertIn(filename, os.listdir(os.path.dirname(os.path.abspath(__file__))))
+
+
+
 def setupTestWavefront():
     """ Utility to setup a Gaussian wavefront. Geometry corresponds to SPB-SFX Day1 configuration. """
 

unittest/Test.py

@@ -0,0 +1,27 @@
+"""
+:module: Top level test module for unittests in WPG.
+"""
+import unittest
+import sys
+
+# Import suites to run.
+from wpg_test.OpticalElementsTest import EmptyTest
+
+# Define the encapsulating test suite.
+def suite():
+    suites = [
+            unittest.makeSuite(EmptyTest, 'test'),
+             ]
+
+    return unittest.TestSuite(suites)
+
+# Run the top level suite and return a success status code. This enables running an automated git-bisect.
+if __name__=="__main__":
+
+    result = unittest.TextTestRunner(verbosity=2).run(suite())
+
+    if result.wasSuccessful():
+        print '---> OK <---'
+        sys.exit(0)
+
+    sys.exit(1)

unittest/wpg_test/OpticalElementsTest.py

@@ -0,0 +1,58 @@
+""" Test module for wpg.OpticalElements.
+
+    :author : Carsten Fortmann-Grote
+    :institution : European XFEL GmbH, Holzkoppel 4, 22869 Schenefeld, Germany
+    :creation date: 20170329
+
+"""
+import os, sys, shutil
+import unittest
+
+
+# Add sources to global namespace.
+sys.path.insert(0,os.path.abspath(os.path.join(os.path.abspath(os.path.dirname(__file__)),'..','..')))
+
+# Import the class to test.
+from wpg import optical_elements
+
+
+class EmptyTest(unittest.TestCase):
+    """
+    Test class for the screen class.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        """ Setting up the test class. """
+        pass
+
+    @classmethod
+    def tearDownClass(cls):
+        """ Tearing down the test class. """
+        pass
+
+    def setUp(self):
+        """ Setting up a test. """
+        self.__files_to_remove = []
+        self.__dirs_to_remove = []
+
+    def tearDown(self):
+        """ Tearing down a test. """
+
+        for f in self.__files_to_remove:
+            if os.path.isfile(f): os.remove(f)
+        for d in self.__dirs_to_remove:
+            if os.path.isdir(d): shutil.rmtree(d)
+
+    def testDefaultConstruction(self):
+        """ Testing the construction of an empty element."""
+        empty = optical_elements.Empty()
+
+        # Check inheritance.
+        self.assertIsInstance(empty, optical_elements.Empty)
+        self.assertIsInstance(empty, optical_elements.WPGOpticalElement)
+        self.assertIsInstance(empty, object)
+
+if __name__ == '__main__':
+    unittest.main()
+