Fix #6003 #5976 wavelength, frequency and amplifier tab refactor (#6076)

sirepo/package_data/static/json/silas-schema.json

@@ -18,6 +18,8 @@
             ["phase", "Phase by Slice"],
             ["photons", "Number of Photons by Slice"],
             ["longitudinal_photons", "Longitudinal Photons"],
+            ["longitudinal_wavelength", "Longitudinal Wavelength"],
+            ["longitudinal_frequency", "Longitudinal Frequency"],
             ["longitudinal_intensity", "Longitudinal Central Intensity"],
             ["total_intensity", "Total Intensity"],
             ["total_phase", "Total Phase"]

sirepo/package_data/template/silas/parameters.py.jinja

@@ -45,6 +45,10 @@ class Watchpoint:
         def ranges(pulse, idx):
             return laser_pulse_ranges(pulse.slice_wfr(idx))
 
+        def longitudinal_wavelength_and_frequency(laser_pulse, slice_index):
+            e = laser_pulse.slice[slice_index].photon_e_ev
+            return (const.h * const.c)/(const.e * e) * 1e9, (2.0 * numpy.pi * e * const.e)/const.h
+
         h5_data = PKDict()
         lp = locals()["pulse"]
         t = lp.extract_total_2d_elec_fields()
@@ -53,6 +57,7 @@ class Watchpoint:
             nx = lp.slice[s].nx_slice
             ny = nx
             i = locals()["intensity"](pulse, s)
+            w, f = longitudinal_wavelength_and_frequency(lp, s)
             h5_data[s] = PKDict(
                 ranges=locals()["ranges"](pulse, s),
                 intensity=i,
@@ -62,6 +67,8 @@ class Watchpoint:
                 longitudinal_intensity=i[int(nx/2)][int(ny/2)],
                 total_intensity=0.5 *const.c *const.epsilon_0 *(t.re**2.0 + t.im**2.0),
                 total_phase=p,
+                longitudinal_wavelength=w,
+                longitudinal_frequency=f,
             )
         template_common.write_dict_to_h5(h5_data, f"results{self.count}.h5")
         self.count += 1

sirepo/template/silas.py

@@ -82,14 +82,14 @@ def _crystal_or_watch(frame_args, element):
     frame_args.sim_in.report = r
     count, element = _report_to_file_index(frame_args.sim_in, r)
     if "beamlineAnimation" in r or r in _SIM_DATA.SOURCE_REPORTS:
-        return _laser_pulse_plot(
-            frame_args.run_dir,
-            _crystal_or_watch(frame_args, element),
-            frame_args.sim_in,
-            count,
-            element,
-            frame_args.frameIndex,
-        )
+        return _LaserPulsePlot(
+            run_dir=frame_args.run_dir,
+            plot_type=_crystal_or_watch(frame_args, element),
+            sim_in=frame_args.sim_in,
+            element_index=count,
+            element=element,
+            slice_index=frame_args.frameIndex,
+        ).gen()
     raise AssertionError("unknown sim_frame report: {}".format(r))
 
 
@@ -418,118 +418,145 @@ def _iterate_beamline(state, data, callback):
         callback(state, e, dz)
 
 
-def _laser_pulse_plot(run_dir, plot_type, sim_in, element_index, element, slice_index):
-    def _cell_volume(element):
-        if _is_crystal(element):
+class _LaserPulsePlot(PKDict):
+    _SCALAR_PLOTS = (
+        "longitudinal_intensity",
+        "longitudinal_frequency",
+        "longitudinal_wavelength",
+    )
+
+    _PLOT_LABELS = PKDict(
+        longitudinal_intensity="Intensity",
+        total_intensity="Total Intensity",
+        total_phase="Total Phase",
+        longitudinal_frequency="Frequency [Rad/s]",
+        longitudinal_wavelength="Wavelength [nm]",
+        longitudinal_photons="Total Number of Photons",
+        excited_states_longitudinal="Excited States",
+        excited_states="Excited States Slice #{slice_index}",
+        phase="Phase Slice #{slice_index}",
+        intensity="Intensity Slice #{slice_index}",
+        photons="Photons Slice #{slice_index}",
+    )
+
+    _X_LABELS = PKDict(
+        excited_states_longitudinal="Crystal Slice",
+        longitudinal_photons="Crystal width [cm]",
+        longitudinal_intensity="Pulse Slice",
+        longitudinal_frequency="Pulse Slice",
+        longitudinal_wavelength="Pulse Slice",
+    )
+
+    _Z_LABELS = PKDict(
+        total_phase="Phase [rad]",
+        total_intensity="",
+        intensity="",
+        phase="Phase [rad]",
+        photons="Photons [1/m³]",
+        excited_states="Number [1/m³]",
+    )
+
+    def _cell_volume(self):
+        if self._is_crystal():
             return (
-                ((2 * element.inversion_mesh_extent) / element.inversion_n_cells) ** 2
-                * element.length
-                / element.nslice
+                (
+                    (2 * self.element.inversion_mesh_extent)
+                    / self.element.inversion_n_cells
+                )
+                ** 2
+                * self.element.length
+                / self.element.nslice
             )
         return None
 
-    def _fname(element):
-        if _is_crystal(element):
+    def _fname(self):
+        if self._is_crystal():
             return _CRYSTAL_FILE
         return _RESULTS_FILE
 
-    def _index(index, plot_type):
-        if plot_type == "longitudinal_photons":
+    def _index(self, index):
+        if self.plot_type == "longitudinal_photons":
             return index
         return index + 1
 
-    def _is_crystal(element):
-        return element and element.type == "crystal"
+    def _is_crystal(self):
+        return self.element and self.element.type == "crystal"
 
-    def _is_longitudinal_plot(plot_type):
-        return "longitudinal" in plot_type
+    def _is_longitudinal_plot(self):
+        return "longitudinal" in self.plot_type
 
-    def _label(plot_type, slice_index):
-        if plot_type == "longitudinal_intensity":
-            return "Intensity"
-        if plot_type == "longitudinal_photons":
-            return "Total Number of Photons"
-        if plot_type == "excited_states_longitudinal":
-            return "Excited States"
-        return _title(plot_type, slice_index)
+    def _plot_label(self):
+        return self._PLOT_LABELS[self.plot_type].format(
+            slice_index=self.slice_index + 1
+        )
 
-    def _nslice(element, file):
-        if _is_crystal(element):
-            return element.nslice
+    def _nslice(self, file):
+        if self._is_crystal():
+            return self.element.nslice
         return len(file)
 
-    def _title(plot_type, slice_index):
-        if plot_type in ("total_intensity", "total_phase"):
-            return plot_type.replace("_", " ").title()
-        return plot_type.replace("_", " ").title() + " Slice #" + str(slice_index + 1)
+    def _x_label(self):
+        return self._X_LABELS[self.plot_type]
 
-    def _x_label(plot_type):
-        return PKDict(
-            excited_states_longitudinal="Crystal Slice",
-            longitudinal_photons="Crystal width [cm]",
-            longitudinal_intensity="Pulse Slice",
-        )[plot_type]
-
-    def _y_value(element, index, file, cell_volume):
-        if _is_crystal(element):
-            return numpy.sum(numpy.array(file[f"{index}/excited_states"]) * cell_volume)
-        y = numpy.array(file[f"{index}/{plot_type}"])
-        if plot_type == "longitudinal_intensity":
+    def _y_value(self, index, file):
+        if self._is_crystal():
+            return numpy.sum(
+                numpy.array(file[f"{index}/excited_states"]) * self._cell_volume()
+            )
+        y = numpy.array(file[f"{index}/{self.plot_type}"])
+        if self.plot_type in self._SCALAR_PLOTS:
             return y
         return numpy.sum(y)
 
-    def _z_label(plot_type):
-        return PKDict(
-            total_phase="Phase [rad]",
-            total_intensity="",
-            intensity="",
-            phase="Phase [rad]",
-            photons="Photons [1/m³]",
-            excited_states="Number [1/m³]",
-        )[plot_type]
-
-    for _ in range(_MAX_H5_READ_TRIES):
-        try:
-            with h5py.File(
-                run_dir.join(_fname(element).format(element_index)), "r"
-            ) as f:
-                if _is_longitudinal_plot(plot_type):
-                    x = []
-                    y = []
-                    nslice = _nslice(element, f)
-                    if element:
-                        element.nslice = nslice
-                    for idx in range(nslice):
-                        x.append(_index(idx, plot_type))
-                        y.append(_y_value(element, idx, f, _cell_volume(element)))
-                    return template_common.parameter_plot(
-                        x,
-                        [
-                            PKDict(
-                                points=y,
-                                label=_label(plot_type, 0),
-                            ),
-                        ],
-                        PKDict(),
-                        PKDict(
-                            x_label=_x_label(plot_type),
-                        ),
+    def _z_label(self):
+        return self._Z_LABELS[self.plot_type]
+
+    def _gen_longitudinal(self, element_file):
+        x = []
+        y = []
+        nslice = self._nslice(element_file)
+        if self.element:
+            self.element.nslice = nslice
+        for idx in range(nslice):
+            x.append(self._index(idx))
+            y.append(self._y_value(idx, element_file))
+        return template_common.parameter_plot(
+            x,
+            [
+                PKDict(
+                    points=y,
+                    label=self._plot_label(),
+                ),
+            ],
+            PKDict(),
+            PKDict(
+                x_label=self._x_label(),
+            ),
+        )
+
+    def gen(self):
+        for _ in range(_MAX_H5_READ_TRIES):
+            try:
+                with h5py.File(
+                    self.run_dir.join(self._fname().format(self.element_index)), "r"
+                ) as f:
+                    if self._is_longitudinal_plot():
+                        return self._gen_longitudinal(f)
+                    d = template_common.h5_to_dict(f, str(self.slice_index))
+                    r = d.ranges
+                    z = d[self.plot_type]
+                    return PKDict(
+                        title=self._plot_label(),
+                        x_range=[r.x[0], r.x[1], len(z)],
+                        y_range=[r.y[0], r.y[1], len(z[0])],
+                        x_label="Horizontal Position [m]",
+                        y_label="Vertical Position [m]",
+                        z_label=self._z_label(),
+                        z_matrix=z,
                     )
-                d = template_common.h5_to_dict(f, str(slice_index))
-                r = d.ranges
-                z = d[plot_type]
-                return PKDict(
-                    title=_title(plot_type, slice_index),
-                    x_range=[r.x[0], r.x[1], len(z)],
-                    y_range=[r.y[0], r.y[1], len(z[0])],
-                    x_label="Horizontal Position [m]",
-                    y_label="Vertical Position [m]",
-                    z_label=_z_label(plot_type),
-                    z_matrix=z,
-                )
-        except BlockingIOError as e:
-            time.sleep(3)
-    raise AssertionError("Report is unavailable")
+            except BlockingIOError as e:
+                time.sleep(3)
+        raise AssertionError("Report is unavailable")
 
 
 def _laser_pulse_report(value_index, filename, title, label):