Skip to content

Commit

Permalink
fix #5998 improve hellweg source generation, fixed energy scaling
Browse files Browse the repository at this point in the history
  • Loading branch information
moellep committed Jun 21, 2023
1 parent 9a92586 commit 23e1527
Show file tree
Hide file tree
Showing 4 changed files with 34 additions and 48 deletions.
17 changes: 13 additions & 4 deletions sirepo/package_data/template/hellweg/parameters.py.jinja
Original file line number Diff line number Diff line change
@@ -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
"""
""")

s = rshellweg.solver.BeamSolver("{{ iniFile }}", "{{ inputFile }}")
s.solve()
s.save_output("{{ outputFile }}")
s.dump_bin("{{ dumpFile }}")
19 changes: 3 additions & 16 deletions sirepo/pkcli/hellweg.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@
from sirepo.template import template_common
import copy
import py.path
import rshellweg.solver
import sirepo.template


def run(cfg_dir):
Expand All @@ -20,13 +18,13 @@ def run(cfg_dir):
Args:
cfg_dir (str): directory to run hellweg in
"""
_run_hellweg(cfg_dir)
template_common.exec_parameters()
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),
Expand All @@ -38,15 +36,4 @@ def run(cfg_dir):


def run_background(cfg_dir):
_run_hellweg(cfg_dir)


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)
s = rshellweg.solver.BeamSolver(t.HELLWEG_INI_FILE, t.HELLWEG_INPUT_FILE)
s.solve()
s.save_output(t.HELLWEG_SUMMARY_FILE)
s.dump_bin(t.HELLWEG_DUMP_FILE)
template_common.exec_parameters()
31 changes: 13 additions & 18 deletions sirepo/template/hellweg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)


Expand Down Expand Up @@ -68,24 +71,12 @@ def background_percent_complete(report, run_dir, is_running):
)


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)
def get_data_file(run_dir, model, frame, options):
return HELLWEG_DUMP_FILE

with open('defaults.ini', 'w') as f:
f.write(ini_file)

s = solver.BeamSolver('defaults.ini', 'input.txt')
s.solve()
s.save_output('output.txt')
""".format(
_generate_parameters_file(data, is_parallel=len(data.models.beamline))
)
def python_source_for_model(data, model, qcall, **kwargs):
return _generate_parameters_file(data, is_parallel=len(data.models.beamline))


def remove_last_frame(run_dir):
Expand Down Expand Up @@ -379,7 +370,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"])
Expand All @@ -389,6 +380,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)


Expand Down
15 changes: 5 additions & 10 deletions sirepo/template/hellweg_dump_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
from pykern.pkdebug import pkdc, pkdexc, pkdlog, pkdp
import ctypes
import math
import numpy

_LIVE_PARTICLE = 0
_LOSS_VALUES = [
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

0 comments on commit 23e1527

Please sign in to comment.