Fix #5996: particles keyword

sirepo/package_data/static/js/hellweg.js

@@ -240,6 +240,9 @@ SIREPO.app.controller('HellwegSourceController', function (appState, panelState,
         panelState.showField('beam', 'longitudinalFile1d', beam.longitudinalDistribution == 'file1d');
         panelState.showField('beam', 'longitudinalFile2d', beam.longitudinalDistribution == 'file2d');
         panelState.showField('beam', 'cstFile', ! isDistribution);
+        ['particleParamA', 'particleParamQ'].forEach(function(f) {
+            panelState.showField('beam', f, beam.particleKeyword === 'ions');
+        });
     }
 
     function updateCurvature() {
@@ -288,7 +291,7 @@ SIREPO.app.controller('HellwegSourceController', function (appState, panelState,
         updateAllFields();
     };
 
-    appState.watchModelFields($scope, ['beam.transversalDistribution', 'beam.longitudinalDistribution', 'beam.spaceCharge', 'beam.beamDefinition', 'sphericalDistribution.curvature', 'sphericalDistribution.curvatureFactor', 'energyPhaseDistribution.distributionType'], updateBeamFields);
+    appState.watchModelFields($scope, ['beam.transversalDistribution', 'beam.longitudinalDistribution', 'beam.spaceCharge', 'beam.beamDefinition', 'sphericalDistribution.curvature', 'sphericalDistribution.curvatureFactor', 'energyPhaseDistribution.distributionType', 'beam.particleKeyword'], updateBeamFields);
     appState.watchModelFields($scope, ['solenoid.sourceDefinition'], updateSolenoidFields);
     appState.whenModelsLoaded($scope, updateAllFields);
 });

sirepo/package_data/static/json/hellweg-schema.json

@@ -69,6 +69,11 @@
             ["x", "Radius (x)"],
             ["y", "Radius (y)"]
         ],
+        "ParticleSpecies": [
+            ["electrons", "Electron"],
+            ["protons", "Proton"],
+            ["ions", "Ion"]
+        ],
         "PlotRangeType": [
             ["none", "None"],
             ["fit", "Fit Data"],
@@ -167,16 +172,19 @@
             "transversalDistribution": ["Transversal Distribution", "TransversalDistribution"],
             "longitudinalDistribution": ["Longitudinal Distribution", "LongitudinalDistribution"],
             "current": ["Current [A]", "Float"],
+            "particleKeyword": ["Particle Species", "ParticleSpecies", "electrons"],
+            "particleParamA": ["A [a.m.u]", "Float", 1],
+            "particleParamQ": ["Q [ee]", "Integer", 1],
             "numberOfParticles": ["Particle Limit", "Integer", 1000, "", 1000],
             "spaceCharge": ["Space Charge Algorithm", "SpaceChargeAlgorithm"],
             "spaceChargeCore": ["Space Charge Ellipsoid Core (rms)", "Float", 3.0],
             "beamDefinition": ["Beam Particle Definition", "BeamDefinition", "transverse_longitude"],
-            "cstCompress": ["Compress Particles into One Bunch", "Boolean"],
-            "transversalFile2d": ["Transversal File (r, r')", "InputFile"],
-            "transversalFile4d": ["Transversal File (x, x', y, y')", "InputFile"],
-            "longitudinalFile1d": ["Longitudinal File (W)", "InputFile"],
-            "longitudinalFile2d": ["Longitudinal File (φ, W)", "InputFile"],
-            "cstFile": ["CST File", "InputFile"]
+            "cstCompress": ["Compress Particles into One Bunch", "Boolean", "0"],
+            "transversalFile2d": ["Transversal File (r, r')", "InputFile", ""],
+            "transversalFile4d": ["Transversal File (x, x', y, y')", "InputFile", ""],
+            "longitudinalFile1d": ["Longitudinal File (W)", "InputFile", ""],
+            "longitudinalFile2d": ["Longitudinal File (φ, W)", "InputFile", ""],
+            "cstFile": ["CST File", "InputFile", ""]
         },
         "beamAnimation": {
             "reportType": ["Report", "BeamReportType"],
@@ -301,6 +309,9 @@
             "title": "Beam",
             "basic": [
                 ["Main", [
+                    "particleKeyword",
+                    "particleParamA",
+                    "particleParamQ",
                     "current",
                     "numberOfParticles",
                     "spaceCharge",

sirepo/package_data/template/hellweg/examples/ions.json

@@ -0,0 +1,128 @@
+{
+    "models": {
+        "beam": {
+            "beamDefinition": "transverse_longitude",
+            "cstCompress": "0",
+            "cstFile": "",
+            "current": 0.01,
+            "longitudinalDistribution": "norm2d",
+            "longitudinalFile1d": "",
+            "longitudinalFile2d": "",
+            "numberOfParticles": 1000,
+            "particleKeyword": "ions",
+            "particleParamA": 12,
+            "particleParamQ": 6,
+            "spaceCharge": "none",
+            "spaceChargeCore": 3,
+            "transversalDistribution": "sph2d",
+            "transversalFile2d": "",
+            "transversalFile4d": ""
+        },
+        "beamAnimation": {
+            "colorMap": "viridis",
+            "framesPerSecond": "5",
+            "histogramBins": 100,
+            "horizontalOffset": 0,
+            "horizontalSize": 0,
+            "isRunning": 0,
+            "notes": "",
+            "plotRangeType": "fit",
+            "plotScale": "linear",
+            "reportType": "x-y",
+            "startTime": 1490386853,
+            "verticalOffset": 0,
+            "verticalSize": 0
+        },
+        "beamHistogramAnimation": {
+            "framesPerSecond": "5",
+            "histogramBins": 100,
+            "notes": "",
+            "reportType": "r",
+            "startTime": 1490386853
+        },
+        "beamHistogramReport": {
+            "histogramBins": 100,
+            "reportType": "r"
+        },
+        "beamReport": {
+            "histogramBins": 100,
+            "reportType": "x-y"
+        },
+        "beamline": [
+            {
+                "frequency": 2856,
+                "id": 2,
+                "inputPower": 20,
+                "phaseShift": 0,
+                "type": "powerElement"
+            },
+            {
+                "acceleratingInvariant": 1160,
+                "aperture": 0.06,
+                "attenuation": 0,
+                "id": 3,
+                "phaseAdvance": 220,
+                "phaseVelocity": 0.3,
+                "repeat": 15,
+                "type": "cellsElement"
+            }
+        ],
+        "ellipticalDistribution": {
+            "aX": 2,
+            "bY": 0.5,
+            "rmsDeviationFactor": 1,
+            "rotationAngle": 0
+        },
+        "energyPhaseDistribution": {
+            "distributionType": "gaussian",
+            "energyDeviation": 0.1,
+            "energySpread": 0.001,
+            "meanEnergy": 41.5,
+            "meanPhase": 30,
+            "phaseDeviation": 5.72957795131,
+            "phaseLength": 1
+        },
+        "parameterAnimation": {
+            "notes": "",
+            "reportType": "wav-wmax"
+        },
+        "particleAnimation": {
+            "notes": "",
+            "renderCount": "300",
+            "reportType": "r"
+        },
+        "simulation": {
+            "folder": "/Examples",
+            "name": "Ions"
+        },
+        "simulationSettings": {
+            "allowBackwardWaves": "0",
+            "meshPoints": 20,
+            "smoothing": 0.95,
+            "splineInterpolation": "0"
+        },
+        "solenoid": {
+            "fieldStrength": 1000,
+            "fringeRegion": 1,
+            "length": 20,
+            "solenoidFile": "",
+            "sourceDefinition": "none",
+            "z0": 0
+        },
+        "sphericalDistribution": {
+            "curvature": "flat",
+            "curvatureFactor": 0,
+            "radialLimit": 0.1,
+            "thermalEmittance": 0
+        },
+        "twissDistribution": {
+            "horizontalAlpha": 0.14,
+            "horizontalBeta": 2,
+            "horizontalEmittance": 0.0005,
+            "verticalAlpha": 0.28,
+            "verticalBeta": 4,
+            "verticalEmittance": 0.0015
+        }
+    },
+    "simulationType": "hellweg"
+}

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

@@ -20,6 +20,7 @@ pykern.pkio.write_text(
 {{ solenoidCommand }}
 {{ chargeCommand }}
 {{ beamCommand }}
+{{ particleSpeciesCommand }}
 {{ currentCommand }}
 {{ latticeCommands }}
 END

sirepo/sim_data/hellweg.py

@@ -19,6 +19,7 @@ def fixup_old_data(cls, data, **kwargs):
         cls._init_models(
             dm,
             (
+                "beam",
                 "beamAnimation",
                 "beamHistogramAnimation",
                 "parameterAnimation",

sirepo/template/hellweg.py

@@ -106,8 +106,8 @@ def sim_frame_beamAnimation(frame_args):
     )
     x, y = frame_args.reportType.split("-")
     values = [
-        hellweg_dump_reader.get_points(beam_info, x),
-        hellweg_dump_reader.get_points(beam_info, y),
+        hellweg_dump_reader.get_points(beam_info, x, data.models.beam.particleKeyword),
+        hellweg_dump_reader.get_points(beam_info, y, data.models.beam.particleKeyword),
     ]
     model["x"] = x
     model["y"] = y
@@ -131,7 +131,9 @@ def sim_frame_beamHistogramAnimation(frame_args):
     beam_info = hellweg_dump_reader.beam_info(
         _dump_file(frame_args.run_dir), frame_args.frameIndex
     )
-    points = hellweg_dump_reader.get_points(beam_info, frame_args.reportType)
+    points = hellweg_dump_reader.get_points(
+        beam_info, frame_args.reportType, frame_args.sim_in.models.beam.particleKeyword
+    )
     hist, edges = numpy.histogram(
         points, template_common.histogram_bins(frame_args.histogramBins)
     )
@@ -181,6 +183,7 @@ def sim_frame_particleAnimation(frame_args):
         _dump_file(frame_args.run_dir),
         frame_args.reportType,
         int(frame_args.renderCount),
+        frame_args.sim_in.models.beam.particleKeyword,
     )
     x = particle_info["z_values"]
     y = particle_info["y_values"]
@@ -226,7 +229,13 @@ def _compute_range_across_files(run_dir, **kwargs):
     for frame in range(beam_header.NPoints):
         beam_info = hellweg_dump_reader.beam_info(dump_file, frame)
         for field in res:
-            values = hellweg_dump_reader.get_points(beam_info, field)
+            values = hellweg_dump_reader.get_points(
+                beam_info,
+                field,
+                simulation_db.read_json(
+                    run_dir.join(template_common.INPUT_BASE_NAME)
+                ).models.beam.particleKeyword,
+            )
             if not values:
                 pass
             elif res[field]:
@@ -296,6 +305,17 @@ def _generate_charge(models):
     )
 
 
+def _generate_particle_species(models):
+    p = models.beam.particleKeyword.upper()
+    if p == "IONS":
+        return "PARTICLES {} {} {}".format(
+            p,
+            models.beam.particleParamA,
+            models.beam.particleParamQ,
+        )
+    return "PARTICLES {}".format(p)
+
+
 def _generate_current(models):
     return "CURRENT {} {}".format(models.beam.current, models.beam.numberOfParticles)
 
@@ -388,6 +408,7 @@ def _generate_parameters_file(data, run_dir=None, is_parallel=False):
     v["beamCommand"] = _generate_beam(data["models"])
     v["currentCommand"] = _generate_current(data["models"])
     v["chargeCommand"] = _generate_charge(data["models"])
+    v["particleSpeciesCommand"] = _generate_particle_species(data["models"])
     if is_parallel:
         v["latticeCommands"] = _generate_lattice(data["models"])
     else:

sirepo/template/hellweg_dump_reader.py

@@ -6,6 +6,7 @@
 """
 
 from pykern.pkdebug import pkdc, pkdexc, pkdlog, pkdp
+from pykern.pkcollections import PKDict
 import ctypes
 import math
 
@@ -52,7 +53,12 @@
     "et",
     "ent",
 ]
-_We0 = 0.5110034e6
+
+_W0_PARTICLE_CONSTANT = PKDict(
+    electrons=0.5110034e6,
+    protons=938.272013e6,
+    ions=931.494028e6,
+)
 
 _BEAM_PARAMETER = {
     # mod() in hellweg means abs()
@@ -72,7 +78,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_ev(p.g),
+    "w": lambda p, lmb: p.g,
 }
 
 _STRUCTURE_PARAMETER = {
@@ -264,22 +270,22 @@ def get_parameter_title(field):
     return _STRUCTURE_TITLE[field]
 
 
-def get_points(info, field):
+def get_points(info, field, particle_species):
     res = []
     fn = _BEAM_PARAMETER[field]
     lmb = info["BeamHeader"].beam_lmb
 
     for p in info["Particles"]:
         if p.lost == _LIVE_PARTICLE:
-            res.append(fn(p, lmb))
+            res.append(_apply_beam_fn(field, fn, p, lmb, particle_species))
     return res
 
 
 def parameter_index(name):
     return _STRUCTURE_VALUES.index(name)
 
 
-def particle_info(filename, field, count):
+def particle_info(filename, field, count, particle_species):
     info = {}
     with open(filename, "rb") as f:
         header = THeader()
@@ -317,7 +323,7 @@ def particle_info(filename, field, count):
                 p = TParticle()
                 assert f.readinto(p) == particle_size
                 if p.lost == _LIVE_PARTICLE:
-                    v = yfn(p, lmb)
+                    v = _apply_beam_fn(field, yfn, p, lmb, particle_species)
                     y_map[idx].append(v)
                     if y_range:
                         if v < y_range[0]:
@@ -338,6 +344,11 @@ def particle_info(filename, field, count):
     return info
 
 
-def _gamma_to_ev(g):
-    # TODO(pjm): when we add species, _We0 will be incorrect
-    return _We0 * (g - 1)
+def _apply_beam_fn(field, fn, p, lmb, species):
+    if field == "w":
+        return _gamma_to_ev(fn(p, lmb), species)
+    return fn(p, lmb)
+
+
+def _gamma_to_ev(g, species):
+    return _W0_PARTICLE_CONSTANT[species] * (g - 1)