address review comments

examples/initialize_from_array/README.rst

@@ -11,15 +11,20 @@ or creating a custom distribution.
 This example includes a set of utilities for transforming the beam to the fixed-s coordinates of ImpactX.
 
 In this example, a custom beam is specified at fixed t, transformed to fixed s, and
-then loaded in ImpactX. The custom beam is a ring in x-y,
+then loaded in ImpactX. 
+The custom beam is a ring in x-y,
 with radius r=2 mm,
 radial width :math:`\sigma_r = 5\ \mu`m;
 Gaussian in :math:`p_x` and :math:`p_y` with momentum width :math:`\sigma_p=10`;
 and chirped in z-pz with bunch length :math:`\sigma_z=1` mm,
 mean energy about 10 GeV, 1% uncorrelated energy spread, and z-pz covariance of -0.18.
 
+In specifying the beam at fixed t and transforming to fixed s,
+it is assumed that the local and global coordinate frames align.
+That is, the beam transverse directions x and y are the global x and y directions
+and the beam z/t direction is the global z direction.
 The transformation utility function reproduces the t-to-s and s-to-t transformations
-done internally in ImpactX.
+done internally in ImpactX given this assumption that the beam and global coordinate systems align.
 These utility functions are provided in the following script:
 
 .. dropdown:: Script ``transformation_utilities.py``

examples/initialize_from_array/run_from_array.py

@@ -1,15 +1,18 @@
-import numpy as np
-
-try:
-    import torch
-
-    on_gpu = torch.cuda.is_available()
-except ImportError:
-    on_gpu = False
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 ImpactX contributors
+# Authors: Ryan Sandberg, Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
 
 import amrex.space3d as amr
-import transformation_utilities as pycoord
 from impactx import ImpactX, elements
+from impactx import Config, ImpactX, elements
+import numpy as np
+import transformation_utilities as pycoord
+
+on_gpu = Config.have_gpu
 
 ################
 
@@ -76,18 +79,18 @@
     dpy_podv = amr.PODVector_real_arena()
     dpt_podv = amr.PODVector_real_arena()
 
-for element in dx:
-    dx_podv.push_back(element)
-for element in dy:
-    dy_podv.push_back(element)
-for element in dt:
-    dt_podv.push_back(element)
-for element in dpx:
-    dpx_podv.push_back(element)
-for element in dpy:
-    dpy_podv.push_back(element)
-for element in dpt:
-    dpt_podv.push_back(element)
+for p_dx in dx:
+    dx_podv.push_back(p_dx)
+for p_dy in dy:
+    dy_podv.push_back(p_dy)
+for p_dt in dt:
+    dt_podv.push_back(p_dt)
+for p_dpx in dpx:
+    dpx_podv.push_back(p_dpx)
+for p_dpy in dpy:
+    dpy_podv.push_back(p_dpy)
+for p_dpt in dpt:
+    dpt_podv.push_back(p_dpt)
 
 pc.add_n_particles(
     dx_podv, dy_podv, dt_podv, dpx_podv, dpy_podv, dpt_podv, qm_eev, bunch_charge_C

examples/initialize_from_array/transformation_utilities.py

@@ -5,7 +5,10 @@ def to_ref_part_t_from_global_t(ref_part, x, y, z, px, py, pz):
     """
     Transform from global coordinates to reference particle coordinates
 
-    This function takes particle bunch data and returns the bunch phase space positions relative to a reference particle
+    This function takes particle bunch data and returns the bunch phase space positions relative to a reference particle.
+    It is assumed that the local and global coordinate frames align.
+    That is, the beam transverse directions x and y are the global x and y directions
+    and the beam z direction is the global z direction.
 
     Parameters
     ---
@@ -42,10 +45,13 @@ def to_ref_part_t_from_global_t(ref_part, x, y, z, px, py, pz):
 
 def to_global_t_from_ref_part_t(ref_part, dx, dy, dz, dpx, dpy, dpz):
     """
-    Transform from reference particle to global coordinates
+    Transform from reference particle to global coordinates.
 
     This function takes particle bunch data relative to a reference particle
-    and returns all particle data in the global coordinate frame
+    and returns all particle data in the global coordinate frame.
+    It is assumed that the local and global coordinate frames align.
+    That is, the beam transverse directions x and y are the global x and y directions
+    and the beam z direction is the global z direction.
 
     Parameters
     ---