Fix #6022 srw scale wavefront during sim

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

@@ -24,6 +24,9 @@ try:
 except:
     pass
 
+{% if in_server %}
+import sirepo.template
+{% endif %}
 import srwl_bl
 import srwlib
 import srwlpy
@@ -431,7 +434,6 @@ def _rsopt_set_params({{ rsOptFuctionSignature() }}):
 
 def epilogue():
     {% if in_server %}
-    import sirepo.template
     sirepo.template.run_epilogue('srw')
     {% else %}
     pass

sirepo/template/srw.py

@@ -11,6 +11,7 @@
 from sirepo import simulation_db
 from sirepo.template import srw_common
 from sirepo.template import template_common
+import srwl_bl
 import array
 import copy
 import math
@@ -23,7 +24,6 @@
 import sirepo.mpi
 import sirepo.sim_data
 import sirepo.util
-import srwl_bl
 import srwlib
 import time
 import traceback
@@ -145,6 +145,8 @@
 
 _JSON_MESSAGE_EXPANSION = 20
 
+_MAX_MESSAGE_BYTES = sirepo.job.cfg().max_message_bytes
+
 _RSOPT_PARAMS = {
     i
     for sublist in [
@@ -438,6 +440,7 @@ def extract_report_data(sim_in):
         ),
     )
     if out.dimensions == 3:
+        res.report = r
         res = _remap_3d(res, allrange, out, dm[r])
     return res
 
@@ -501,28 +504,14 @@ def sim_frame(frame_args):
         _copy_frame_args_into_model(frame_args, r)
     elif "beamlineAnimation" in r:
         wid = int(re.search(r".*?(\d+)$", r)[1])
-        fn = _wavefront_pickle_filename(wid)
-        with open(fn, "rb") as f:
-            wfr = pickle.load(f)
         m = _copy_frame_args_into_model(frame_args, "watchpointReport")
         if wid:
             m.id = wid
             frame_args.sim_in.report = "beamlineAnimation"
             frame_args.sim_in.models.beamlineAnimation = m
-            data_file = _OUTPUT_FOR_MODEL.beamlineAnimation.filename.format(
-                watchpoint_id=wid
-            )
         else:
             frame_args.sim_in.report = "initialIntensityReport"
             frame_args.sim_in.models.initialIntensityReport = m
-            data_file = _OUTPUT_FOR_MODEL.initialIntensityReport.filename
-        srwl_bl.SRWLBeamline().calc_int_from_wfr(
-            wfr,
-            _pol=int(frame_args.polarization),
-            _int_type=int(frame_args.characteristic),
-            _fname=data_file,
-            _pr=False,
-        )
     if "beamlineAnimation" not in r:
         # some reports may be written at the same time as the reader
         # if the file is invalid, wait a bit and try again
@@ -787,6 +776,50 @@ def process_undulator_definition(model):
         return model
 
 
+def process_watch(wid=0):
+    def _op():
+        sim_in = simulation_db.read_json(template_common.INPUT_BASE_NAME)
+        report = sim_in.models[f"beamlineAnimation{wid}"]
+        with open(_wavefront_pickle_filename(wid), "rb") as f:
+            wfr = pickle.load(f)
+        data, mesh = srwl_bl.SRWLBeamline().calc_int_from_wfr(
+            wfr,
+            _pol=int(report.get("polarization", "6")),
+            _int_type=int(report.get("characteristic", "0")),
+            _pr=False,
+        )
+        d, x_range, y_range = _reshape_3d(
+            np.array(data),
+            [0, 0, 0, mesh.xStart, mesh.xFin, mesh.nx, mesh.yStart, mesh.yFin, mesh.ny],
+            report,
+        )
+        new_mesh = srwlib.SRWLRadMesh(
+            _eStart=mesh.eStart,
+            _eFin=mesh.eFin,
+            _ne=mesh.ne,
+            _xStart=x_range[0],
+            _xFin=x_range[1],
+            _nx=x_range[2],
+            _yStart=y_range[0],
+            _yFin=y_range[1],
+            _ny=y_range[2],
+            _zStart=mesh.zStart,
+            _nvx=mesh.nvx,
+            _nvy=mesh.nvy,
+            _nvz=mesh.nvz,
+            _hvx=mesh.hvx,
+            _hvy=mesh.hvy,
+            _hvz=mesh.hvz,
+            _arSurf=mesh.arSurf,
+        )
+        dst = _wavefront_intensity_filename(wid)
+        src = f"tmp_{dst}"
+        srwlib.srwl_uti_save_intens_ascii(d.flatten().tolist(), new_mesh, src)
+        pkio.py_path(src).rename(dst)
+
+    sirepo.mpi.restrict_op_to_first_rank(_op)
+
+
 def python_source_for_model(data, model, qcall, plot_reports=True, **kwargs):
     data.report = model or _SIM_DATA.SRW_RUN_ALL_MODEL
     data.report = re.sub("beamlineAnimation0", "initialIntensityReport", data.report)
@@ -1143,7 +1176,7 @@ def _beamline_animation_percent_complete(run_dir, res):
     res.outputInfo = [
         PKDict(
             modelKey="beamlineAnimation0",
-            filename=_wavefront_pickle_filename(0),
+            filename=_wavefront_intensity_filename(0),
             id=0,
         ),
     ]
@@ -1156,16 +1189,14 @@ def _beamline_animation_percent_complete(run_dir, res):
                 PKDict(
                     waitForData=True,
                     modelKey=f"beamlineAnimation{item.id}",
-                    filename=_wavefront_pickle_filename(item.id),
+                    filename=_wavefront_intensity_filename(item.id),
                     id=item.id,
                 )
             )
     count = 0
     for info in res.outputInfo:
         try:
-            with open(info.filename, "rb") as f:
-                # TODO(pjm): instead look at last byte == pickle.STOP, see template_common.read_last_csv_line()
-                wfr = pickle.load(f)
+            if template_common.read_last_csv_line(pkio.py_path(info.filename)):
                 count += 1
                 info.waitForData = False
         except Exception as e:
@@ -1956,6 +1987,7 @@ def _generate_srw_main(data, plot_reports, beamline_info):
     source_type = data.models.simulation.sourceType
     run_all = report == _SIM_DATA.SRW_RUN_ALL_MODEL or is_for_rsopt
     vp_var = "vp" if is_for_rsopt else "varParam"
+    final_watch = None
     content = [
         f"v = srwl_bl.srwl_uti_parse_options(srwl_bl.srwl_uti_ext_options({vp_var}), use_sys_argv={plot_reports})",
     ]
@@ -1975,10 +2007,12 @@ def _generate_srw_main(data, plot_reports, beamline_info):
             content.append("op = None")
         content.append("v.ws_fne = '{}'".format(_wavefront_pickle_filename(0)))
         prev_wavefront = None
+        prev_watch = 0
         names = []
         for n in beamline_info.names:
             names.append(n)
             if n in beamline_info.watches:
+                final_watch = n
                 is_last_watch = n == beamline_info.names[-1]
                 content.append("names = ['" + "','".join(names) + "']")
                 names = []
@@ -1989,6 +2023,11 @@ def _generate_srw_main(data, plot_reports, beamline_info):
                 content.append("op = set_optics(v, names, {})".format(is_last_watch))
                 if not is_last_watch:
                     content.append("srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)")
+                    content.append(
+                        f"sirepo.template.import_module('srw').process_watch(wid={prev_watch})"
+                    )
+                prev_watch = beamline_info.watches[n]
+
     elif run_all or (
         _SIM_DATA.srw_is_beamline_report(report) and len(data.models.beamline)
     ):
@@ -2057,6 +2096,10 @@ def _generate_srw_main(data, plot_reports, beamline_info):
             if plot_reports:
                 content.append("v.tr_pl = 'xz'")
     content.append("srwl_bl.SRWLBeamline(_name=v.name).calc_all(v, op)")
+    if report == "beamlineAnimation":
+        content.append(
+            f"sirepo.template.import_module('srw').process_watch(wid={beamline_info.watches.get(final_watch, 0)})"
+        )
     return "\n".join(
         [f"    {x}" for x in content] + [""] + ([] if is_for_rsopt else ["main()", ""])
     )
@@ -2162,21 +2205,10 @@ def _process_rsopt_elements(els):
 
 
 def _remap_3d(info, allrange, out, report):
-    x_range = [allrange[3], allrange[4], allrange[5]]
-    y_range = [allrange[6], allrange[7], allrange[8]]
-    ar2d = info.points
-    totLen = int(x_range[2] * y_range[2])
-    n = len(ar2d) if totLen > len(ar2d) else totLen
-    ar2d = np.reshape(ar2d[0:n], (int(y_range[2]), int(x_range[2])))
-
-    if report.get("usePlotRange", "0") == "1":
-        ar2d, x_range, y_range = _update_report_range(report, ar2d, x_range, y_range)
-    if report.get("useIntensityLimits", "0") == "1":
-        ar2d[ar2d < report.minIntensityLimit] = report.minIntensityLimit
-        ar2d[ar2d > report.maxIntensityLimit] = report.maxIntensityLimit
-    ar2d, x_range, y_range = _resize_report(report, ar2d, x_range, y_range)
-    if report.get("rotateAngle", 0):
-        ar2d, x_range, y_range = _rotate_report(report, ar2d, x_range, y_range, info)
+    ar2d, x_range, y_range = _reshape_3d(np.array(info.points), allrange, report)
+    rotate_angle = report.get("rotateAngle", 0)
+    if rotate_angle and info.title != "Power Density":
+        info.subtitle = info.subtitle + " Image Rotate {}^0".format(rotate_angle)
     if out.units[2]:
         out.labels[2] = "{} [{}]".format(out.labels[2], out.units[2])
     if report.get("useIntensityLimits", "0") == "1":
@@ -2197,16 +2229,33 @@ def _remap_3d(info, allrange, out, report):
     )
 
 
+def _reshape_3d(ar1d, allrange, report):
+    x_range = [allrange[3], allrange[4], allrange[5]]
+    y_range = [allrange[6], allrange[7], allrange[8]]
+    totLen = int(x_range[2] * y_range[2])
+    n = len(ar1d) if totLen > len(ar1d) else totLen
+    ar2d = np.reshape(ar1d[0:n], (int(y_range[2]), int(x_range[2])))
+    if report.get("usePlotRange", "0") == "1":
+        ar2d, x_range, y_range = _update_report_range(report, ar2d, x_range, y_range)
+    if report.get("useIntensityLimits", "0") == "1":
+        ar2d[ar2d < report["minIntensityLimit"]] = report["minIntensityLimit"]
+        ar2d[ar2d > report["maxIntensityLimit"]] = report["maxIntensityLimit"]
+    ar2d, x_range, y_range = _resize_report(report, ar2d, x_range, y_range)
+    if report.get("rotateAngle", 0):
+        ar2d, x_range, y_range = _rotate_report(report, ar2d, x_range, y_range)
+    return ar2d, x_range, y_range
+
+
 def _resize_report(report, ar2d, x_range, y_range):
+    from pykern.pkdebug import pkdc
+
     width_pixels = int(report.get("intensityPlotsWidth", 0))
     if not width_pixels:
         # upper limit is browser's max html canvas size
         width_pixels = _CANVAS_MAX_SIZE
     # roughly 20x size increase for json
-    if ar2d.size * _JSON_MESSAGE_EXPANSION > sirepo.job.cfg().max_message_bytes:
-        max_width = int(
-            math.sqrt(sirepo.job.cfg().max_message_bytes / _JSON_MESSAGE_EXPANSION)
-        )
+    if ar2d.size * _JSON_MESSAGE_EXPANSION > _MAX_MESSAGE_BYTES:
+        max_width = int(math.sqrt(_MAX_MESSAGE_BYTES / _JSON_MESSAGE_EXPANSION))
         if max_width < width_pixels:
             pkdc(
                 "auto scaling dimensions to fit message size. size: {}, max_width: {}",
@@ -2238,8 +2287,9 @@ def _resize_report(report, ar2d, x_range, y_range):
     return ar2d, x_range, y_range
 
 
-def _rotate_report(report, ar2d, x_range, y_range, info):
+def _rotate_report(report, ar2d, x_range, y_range):
     from scipy import ndimage
+    from pykern.pkdebug import pkdc
 
     rotate_angle = report.rotateAngle
     rotate_reshape = report.get("rotateReshape", "0") == "1"
@@ -2272,8 +2322,8 @@ def _rotate_report(report, ar2d, x_range, y_range, info):
 
     x_range[2] = ar2d.shape[1]
     y_range[2] = ar2d.shape[0]
-    if info.title != "Power Density":
-        info.subtitle = info.subtitle + " Image Rotate {}^0".format(rotate_angle)
+    # if info.title != "Power Density":
+    #    info.subtitle = info.subtitle + " Image Rotate {}^0".format(rotate_angle)
     return ar2d, x_range, y_range
 
 
@@ -2642,6 +2692,14 @@ def file_attrs_ok(attrs):
         )
 
 
+def _wavefront_intensity_filename(el_id):
+    return (
+        _OUTPUT_FOR_MODEL.beamlineAnimation.filename.format(watchpoint_id=el_id)
+        if el_id
+        else _OUTPUT_FOR_MODEL.initialIntensityReport.filename
+    )
+
+
 def _wavefront_pickle_filename(el_id):
     if el_id:
         return f"wid-{el_id}.pkl"

tests/template/srw_resize_3d_test.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+"""PyTest for :mod:`sirepo.template.srw`
+
+:copyright: Copyright (c) 2023 RadiaSoft LLC.  All Rights Reserved.
+:license: http://www.apache.org/licenses/LICENSE-2.0.html
+"""
+from pykern.pkcollections import PKDict
+from pykern.pkunit import pkeq, pkok
+
+_N_BINS = 100
+
+
+def setup_module(module):
+    import os
+
+    # Set max to something reasonable for testing
+    os.environ.update(
+        SIREPO_JOB_MAX_MESSAGE_BYTES=str(_N_BINS * _N_BINS),
+    )
+
+
+def test_srw_resize_3d(fc):
+    from sirepo.template import srw
+    import numpy
+
+    a, xr, yr = srw._reshape_3d(
+        numpy.random.rand(srw._MAX_MESSAGE_BYTES),
+        [0, 0, 0, 0.0, 1.0, _N_BINS, 0.0, 1.0, _N_BINS],
+        PKDict(),
+    )
+    pkok(
+        xr[2] < _N_BINS and yr[2] < _N_BINS,
+        "did not reduce bins nx={} ny={}",
+        xr[2],
+        yr[2],
+    )