Working fix to SASE spectrum

twguest · May 6, 2024 · daef6bf · daef6bf
1 parent b16a81a
commit daef6bf
Show file tree

Hide file tree

Showing 6 changed files with 118 additions and 35 deletions.
diff --git a/docs/source/notebooks/phenom_to_wpg.ipynb b/docs/source/notebooks/phenom_to_wpg.ipynb
diff --git a/docs/source/notebooks/sase_model_pt2.ipynb b/docs/source/notebooks/sase_model_pt2.ipynb
diff --git a/docs/source/notebooks/sase_model_pt3.ipynb b/docs/source/notebooks/sase_model_pt3.ipynb
diff --git a/phenom/source.py b/phenom/source.py
@@ -33,6 +33,7 @@ def sase_pulse(
     t0,
     theta_x,
     theta_y,
+    domain = 'freq'
 ):
     """
     SASE pulse model
@@ -50,15 +51,16 @@ def sase_pulse(
     :param t0: temporal jitter (float)
     :param theta_x: horizontal pointing angle (float)
     :param theta_y: vertical pointing angle (float)
-
+    :param domain: 'time' or 'freq', determines the output of sase spectrum
     """
 
     tfield = linear_SASE_spectrum(
         t,
         pulse_duration=pulse_duration,
         photon_energy=photon_energy,
         bandwidth=bandwidth,
-        t0=t0,)
+        t0=t0,
+        output = domain)
 
     sfield = complex_gaussian_beam(
         x=x,
@@ -116,26 +118,29 @@ def sort_arg_types(args, __type__):
 
     __arrays__ = []
     __lams__ = []
-
+    __strs__ = []
     N = 1  ### number of pulses
 
     for key in __keys__:
         # print(key,type(args[key])) # dev @author: twguest - 13/06/23
-
+        
         if type(args[key]) == np.ndarray:
             __arrays__.append(key)
 
         elif type(args[key]) == FunctionType:
             __lams__.append(key)
-
+
+        elif type(args[key]) == str:
+            __strs__.append(key)
+
     __len__ = [args[key].shape[0] for key in __arrays__]
 
     if len(__arrays__) > 0:
         assert np.all(__len__ == np.mean(__len__)), "all length of all lists and 1d numpy arrays must be equal"
 
         N = int(np.mean(__len__))
 
-    __floats__ = list(set(__keys__) - set(__arrays__) - set(__lams__))
+    __floats__ = list(set(__keys__) - set(__arrays__) - set(__lams__) - set(__strs__))
 
     __set__ = {}
 
@@ -344,6 +349,7 @@ def __init__(
         t0,
         theta_x,
         theta_y,
+        domain = 'freq'
     ):
         """
         initialisation function.
@@ -381,6 +387,7 @@ def __init__(
             t0=t0,
             theta_x=theta_x,
             theta_y=theta_y,
+            domain = domain
         )
 
     def generate_sase_field(self, params):

diff --git a/phenom/spectrum.py b/phenom/spectrum.py
@@ -70,7 +70,7 @@ def linear_SASE_spectrum(t, pulse_duration, photon_energy, bandwidth=1e-04, t0=0
     if output == 'time':
         complex_spectrum = time_domain_field
     elif output == 'freq':
-        complex_spectrum = np.fft.ifftshift(freq_domain_field)
+        complex_spectrum = freq_domain_field
 
 
     return complex_spectrum

diff --git a/phenom/wpg.py b/phenom/wpg.py
@@ -65,9 +65,11 @@ def wpg_converter(save_loc, key):
 
     wfr.params.Mesh.sliceMin = np.min(t)
     wfr.params.Mesh.sliceMax = np.max(t)
-
+    
     wfr.params.wDomain = 'time'
-
+
+    srwlpy.SetRepresElecField(wfr._srwl_wf, 'frequency')
+
     wfr.params.Mesh.xMin = np.min(x)
     wfr.params.Mesh.xMax = np.max(x)
     wfr.params.Mesh.yMin = np.min(y)
@@ -77,5 +79,5 @@ def wpg_converter(save_loc, key):
     wfr.params.Ry = 1
 
     wfr.data.arrEhor = complex_to_wpg(efield)
-    srwlpy.SetRepresElecField(wfr._srwl_wf, 'frequency')
+
     return wfr