Skip to content

Commit

Permalink
for #6022 rework maximum report dimension calculations
Browse files Browse the repository at this point in the history 
- clean up preprocess files after execution
  • Loading branch information
@moellep
moellep committed Sep 16, 2023
1 parent d0d528f commit 1c2841f
Showing 1 changed file with 80 additions and 75 deletions.
155 changes: 80 additions & 75 deletions sirepo/template/srw.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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,
),
]
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 1c2841f

Please sign in to comment.