fix #5998 improve hellweg source generation, fixed energy scaling (#5999

)

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

@@ -1,18 +1,27 @@
 # {{simulation_name}}
 
-ini_file = """
+import pykern.pkio
+import rshellweg.solver
+
+pykern.pkio.write_text(
+    "{{ iniFile }}",
+    """
 [NUMERIC]
 Number of Mesh Points = {{simulationSettings_meshPoints}}
 Spline Interpolation = {{simulationSettings_splineInterpolation}}
 Smoothing = {{simulationSettings_smoothing}}
-"""
+""",
+)
 
-input_file = """
+pykern.pkio.write_text(
+    "{{ inputFile }}",
+    """
 {{ optionsCommand }}
 {{ solenoidCommand }}
 {{ chargeCommand }}
 {{ beamCommand }}
 {{ currentCommand }}
 {{ latticeCommands }}
 END
-"""
+""",
+)

sirepo/pkcli/hellweg.py

@@ -11,7 +11,7 @@
 import copy
 import py.path
 import rshellweg.solver
-import sirepo.template
+import sirepo.template.hellweg
 
 
 def run(cfg_dir):
@@ -20,13 +20,13 @@ def run(cfg_dir):
     Args:
         cfg_dir (str): directory to run hellweg in
     """
-    _run_hellweg(cfg_dir)
+    _run_hellweg()
     sim_in = simulation_db.read_json(template_common.INPUT_BASE_NAME)
     r = sim_in.report
     template_common.write_sequential_result(
         template_common.sim_frame_dispatch(
             copy.deepcopy(sim_in.models[r]).pkupdate(
-                # TODO(e-carlin): Is this right? What should frameIndex be?
+                # show very first frame for source reports
                 frameIndex=0,
                 frameReport=r.replace("Report", "Animation"),
                 run_dir=pkio.py_path(cfg_dir),
@@ -38,14 +38,12 @@ def run(cfg_dir):
 
 
 def run_background(cfg_dir):
-    _run_hellweg(cfg_dir)
+    _run_hellweg()
 
 
-def _run_hellweg(cfg_dir):
-    t = sirepo.template.import_module("hellweg")
-    r = template_common.exec_parameters()
-    pkio.write_text(t.HELLWEG_INPUT_FILE, r.input_file)
-    pkio.write_text(t.HELLWEG_INI_FILE, r.ini_file)
+def _run_hellweg():
+    t = sirepo.template.hellweg
+    template_common.exec_parameters()
     s = rshellweg.solver.BeamSolver(t.HELLWEG_INI_FILE, t.HELLWEG_INPUT_FILE)
     s.solve()
     s.save_output(t.HELLWEG_SUMMARY_FILE)

sirepo/template/hellweg.py

@@ -29,12 +29,15 @@
 WANT_BROWSER_FRAME_CACHE = True
 
 # lattice element is required so make it very short and wide drift
-_DEFAULT_DRIFT_ELEMENT = "DRIFT 1e-16 1e+16 2" + "\n"
+# needs to be 3 cells or else the Space Charge example doesn't produce correct results
+_DEFAULT_DRIFT_ELEMENT = "DRIFT 1e-16 1e+16 3" + "\n"
 
 _HELLWEG_PARSED_FILE = "PARSED.TXT"
 
 _PARAMETER_SCALE = PKDict(
     rb=2.0,
+    wav=1e6,
+    wmax=1e6,
 )
 
 
@@ -68,23 +71,23 @@ def background_percent_complete(report, run_dir, is_running):
     )
 
 
+def get_data_file(run_dir, model, frame, options):
+    return HELLWEG_DUMP_FILE
+
+
 def python_source_for_model(data, model, qcall, **kwargs):
     return """
-from rshellweg import solver
-
 {}
-
-with open('input.txt', 'w') as f:
-    f.write(input_file)
-
-with open('defaults.ini', 'w') as f:
-    f.write(ini_file)
-
-s = solver.BeamSolver('defaults.ini', 'input.txt')
+s = rshellweg.solver.BeamSolver("{}", "{}")
 s.solve()
-s.save_output('output.txt')
-    """.format(
-        _generate_parameters_file(data, is_parallel=len(data.models.beamline))
+s.save_output("{}")
+s.dump_bin("{}")
+""".format(
+        _generate_parameters_file(data, is_parallel=len(data.models.beamline)),
+        HELLWEG_INI_FILE,
+        HELLWEG_INPUT_FILE,
+        HELLWEG_SUMMARY_FILE,
+        HELLWEG_DUMP_FILE,
     )
 
 
@@ -379,7 +382,7 @@ def _generate_options(models):
 
 def _generate_parameters_file(data, run_dir=None, is_parallel=False):
     template_common.validate_models(data, SCHEMA)
-    v = template_common.flatten_data(data["models"], {})
+    v = template_common.flatten_data(data["models"], PKDict())
     v["optionsCommand"] = _generate_options(data["models"])
     v["solenoidCommand"] = _generate_solenoid(data["models"])
     v["beamCommand"] = _generate_beam(data["models"])
@@ -389,6 +392,10 @@ def _generate_parameters_file(data, run_dir=None, is_parallel=False):
         v["latticeCommands"] = _generate_lattice(data["models"])
     else:
         v["latticeCommands"] = _DEFAULT_DRIFT_ELEMENT
+    v.iniFile = HELLWEG_INI_FILE
+    v.inputFile = HELLWEG_INPUT_FILE
+    v.outputFile = HELLWEG_SUMMARY_FILE
+    v.dumpFile = HELLWEG_DUMP_FILE
     return template_common.render_jinja(SIM_TYPE, v)
 
 

sirepo/template/hellweg_dump_reader.py

@@ -8,7 +8,6 @@
 from pykern.pkdebug import pkdc, pkdexc, pkdlog, pkdp
 import ctypes
 import math
-import numpy
 
 _LIVE_PARTICLE = 0
 _LOSS_VALUES = [
@@ -73,7 +72,7 @@
     "phi": lambda p, lmb: p.phi * 180.0 / math.pi,
     "zrel": lambda p, lmb: lmb * p.phi / (2 * math.pi),
     "z0": lambda p, lmb: p.z,
-    "w": lambda p, lmb: _gamma_to_mev(p.g),
+    "w": lambda p, lmb: _gamma_to_ev(p.g),
 }
 
 _STRUCTURE_PARAMETER = {
@@ -131,8 +130,6 @@
     "w": "W [eV]",
 }
 
-_PARTICLE_Y_SCALE_FACTOR = 1e6
-
 # some values are pointers which would never serialize correctly
 # the same size in bytes as a long integer
 _POINTER_TYPE = ctypes.c_long
@@ -335,14 +332,12 @@ def particle_info(filename, field, count):
         assert f.readinto(header) == 0
         y_values = []
         for idx in sorted(y_map.keys()):
-            y_values.append(
-                (numpy.array(y_map[idx]) * _PARTICLE_Y_SCALE_FACTOR).tolist()
-            )
+            y_values.append(y_map[idx])
         info["y_values"] = y_values
-        info["y_range"] = (numpy.array(y_range) * _PARTICLE_Y_SCALE_FACTOR).tolist()
+        info["y_range"] = y_range
     return info
 
 
-def _gamma_to_mev(g):
+def _gamma_to_ev(g):
     # TODO(pjm): when we add species, _We0 will be incorrect
-    return _We0 * (g - 1) * 1e-6
+    return _We0 * (g - 1)