From 1c2841fe24fc6d83ab16a74a15a893f7d3b6b05f Mon Sep 17 00:00:00 2001 From: moellep Date: Sat, 16 Sep 2023 01:01:43 +0000 Subject: [PATCH] for #6022 rework maximum report dimension calculations - clean up preprocess files after execution --- sirepo/template/srw.py | 155 +++++++++++++++++++++-------------------- 1 file changed, 80 insertions(+), 75 deletions(-) diff --git a/sirepo/template/srw.py b/sirepo/template/srw.py index 8bbd21dae3..00efc96e06 100644 --- a/sirepo/template/srw.py +++ b/sirepo/template/srw.py @@ -11,19 +11,19 @@ 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 glob import math import numpy as np import os import pickle import pykern.pkjson import re -import sirepo.job import sirepo.mpi import sirepo.sim_data import sirepo.util +import srwl_bl import srwlib import srwlpy import time @@ -37,6 +37,8 @@ PARSED_DATA_ATTR = "srwParsedData" _CANVAS_MAX_SIZE = 65535 +_LOG_DIR = "__srwl_logs__" +_MAX_REPORT_POINTS = 20000000 _MIN_CORES = 3 _OUTPUT_FOR_MODEL = PKDict( @@ -142,11 +144,7 @@ ) _OUTPUT_FOR_MODEL.sourceIntensityReport.title = "E={sourcePhotonEnergy} eV" -_LOG_DIR = "__srwl_logs__" - -_JSON_MESSAGE_EXPANSION = 20 - -_MAX_MESSAGE_BYTES = sirepo.job.cfg().max_message_bytes +_PREPROCESS_PREFIX = "preproc-" _RSOPT_PARAMS = { i @@ -673,6 +671,8 @@ def post_execution_processing( success_exit, **kwargs, ): + for f in glob.glob(str(run_dir.join(_PREPROCESS_PREFIX + "*"))): + os.remove(f) if success_exit: if _SIM_DATA.is_for_ml(compute_model): f = _SIM_DATA.ML_OUTPUT @@ -792,55 +792,70 @@ def process_undulator_definition(model): def process_watch(wid=0): + def _resize_wavefront(wfr): + mesh = wfr.mesh + nx = min(mesh.nx, _CANVAS_MAX_SIZE) + ny = min(mesh.ny, _CANVAS_MAX_SIZE) + _MIN_DIMENSION = int(_MAX_REPORT_POINTS / _CANVAS_MAX_SIZE) + 1 + r = math.sqrt(_MAX_REPORT_POINTS / (nx * ny)) + if nx < _MIN_DIMENSION: + ny = int(_MAX_REPORT_POINTS / nx) + elif r * nx < _MIN_DIMENSION: + nx = _MIN_DIMENSION + ny = int(_MAX_REPORT_POINTS / nx) + elif ny < _MIN_DIMENSION: + nx = int(_MIN_DIMENSION / ny) + elif r * ny < _MIN_DIMENSION: + ny = _MIN_DIMENSION + nx = int(_MAX_REPORT_POINTS / ny) + else: + nx = int(r * nx) + ny = int(r * ny) + pkdc("resized mesh: {}x{}", nx, ny) + # resize the electic fields in the wavefront mesh - note it modifies wfr + srwlpy.ResizeElecFieldMesh( + wfr, + srwlib.SRWLRadMesh( + _eStart=mesh.eStart, + _eFin=mesh.eFin, + _ne=mesh.ne, + _xStart=mesh.xStart, + _xFin=mesh.xFin, + _nx=nx, + _yStart=mesh.yStart, + _yFin=mesh.yFin, + _ny=ny, + _zStart=mesh.zStart, + _nvx=mesh.nvx, + _nvy=mesh.nvy, + _nvz=mesh.nvz, + _hvx=mesh.hvx, + _hvy=mesh.hvy, + _hvz=mesh.hvz, + _arSurf=mesh.arSurf, + ), + [0, 1], + ) + return wfr + def _op(): sim_in = simulation_db.read_json(template_common.INPUT_BASE_NAME) report = sim_in.models[f"beamlineAnimation{wid}"] p = _wavefront_pickle_filename(wid) with open(p, "rb") as f: wfr = pickle.load(f) - # calculate but do not save to file - 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, - ) - pkio.unchecked_remove(p) - - # for the purposes of resizing the wavefront, ignore usePlotRange, useIntensityLimits, and rotateAngle - _, 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, - ignore_plot_params=True, - ) - - # reshaped mesh - 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, - ) - - # resize the electic fields in the wavefront mesh - note it modifies wfr - srwlpy.ResizeElecFieldMesh(wfr, new_mesh, [0, 1]) - - with open(_wavefront_pickle_filename(wid, is_processed=True), "wb") as f: - pickle.dump(wfr, f) + pkdc("original mesh: {}x{}", wfr.mesh.nx, wfr.mesh.ny) + if ( + wfr.mesh.nx < _CANVAS_MAX_SIZE + and wfr.mesh.ny < _CANVAS_MAX_SIZE + and wfr.mesh.nx * wfr.mesh.ny <= _MAX_REPORT_POINTS + ): + pkio.py_path(p).copy( + pkio.py_path(_wavefront_pickle_filename(wid, is_processed=True)) + ) + else: + with open(_wavefront_pickle_filename(wid, is_processed=True), "wb") as f: + pickle.dump(_resize_wavefront(wfr), f) sirepo.mpi.restrict_op_to_first_rank(_op) @@ -1201,7 +1216,6 @@ def _beamline_animation_percent_complete(run_dir, res): res.outputInfo = [ PKDict( modelKey="beamlineAnimation0", - filename=_wavefront_pickle_filename(0, is_processed=True), id=0, ), ] @@ -1214,14 +1228,16 @@ def _beamline_animation_percent_complete(run_dir, res): PKDict( waitForData=True, modelKey=f"beamlineAnimation{item.id}", - filename=_wavefront_pickle_filename(item.id, is_processed=True), id=item.id, ) ) count = 0 for info in res.outputInfo: try: - with open(pkio.py_path(info.filename), "rb") as f: + with open( + pkio.py_path(_wavefront_pickle_filename(info.id, is_processed=True)), + "rb", + ) as f: f.seek(-1, os.SEEK_END) if f.read(1) == pickle.STOP: count += 1 @@ -2259,42 +2275,30 @@ def _remap_3d(info, allrange, out, report): ) -def _reshape_3d(ar1d, allrange, report, ignore_plot_params=False): +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 not ignore_plot_params and report.get("usePlotRange", "0") == "1": + if report.get("usePlotRange", "0") == "1": ar2d, x_range, y_range = _update_report_range(report, ar2d, x_range, y_range) - if not ignore_plot_params and report.get("useIntensityLimits", "0") == "1": + 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, ignore_plot_params=ignore_plot_params - ) - if not ignore_plot_params and report.get("rotateAngle", 0): + 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, ignore_plot_params=False): +def _resize_report(report, ar2d, x_range, y_range): from pykern.pkdebug import pkdc - width_pixels = int(not ignore_plot_params and report.get("intensityPlotsWidth", 0)) + 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 > _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: {}", - ar2d.size, - max_width, - ) - width_pixels = max_width # rescale width and height to maximum of width_pixels if width_pixels and (width_pixels < x_range[2] or width_pixels < y_range[2]): from scipy import ndimage @@ -2731,9 +2735,10 @@ def _wavefront_intensity_filename(el_id, is_processed=True): def _wavefront_pickle_filename(el_id, is_processed=False): - if el_id: - return f"wid-{el_id}{'' if is_processed else '-preproc'}.pkl" - return "initial.pkl" + f = f"wid-{el_id}" if el_id else "initial" + if not is_processed: + f = _PREPROCESS_PREFIX + f + return f"{f}.pkl" def _write_rsopt_files(data, run_dir, ctx):