Merge branch 'master' into 395-screen-image-orientation-is-incorrect

Author: jank324

CHANGELOG.md

@@ -10,7 +10,7 @@
 ### 🚀 Features
 
 - Implement `split` method for the `Solenoid` element (see #380) (@cr-xu)
-- Implement a more robust RPN parser, fixing a bug where short strings in an Elegant variable definition would cause parsing to fail. (see #387) (@amylizzle, @Hespe, @jank324)
+- Implement a more robust RPN parser, fixing a bug where short strings in an Elegant variable definition would cause parsing to fail. (see #387, #417) (@amylizzle, @Hespe, @jank324)
 - Add a `Sextupole` element (see #406) (@jank324, @Hespe)
 
 ### 🐛 Bug fixes
@@ -21,6 +21,7 @@
 - Fix issue where `Quadrupole.tracking_method` was not preserved on cloning (see #404) (@RemiLehe, @jank324)
 - The vertical screen misalignment is now correctly applied to `y` instead of `px` (see #405) (@RemiLehe)
 - Fix issues when generating screen images caused by the sign of particle charges (see #394) (@Hespe, @jank324)
+- Fix an issue where newer versions of `torch` only accept a `torch.Tensor` as input to `torch.rad2deg` (see #417) (@jank324)
 - Fix bug that caused correlations to be lost in the conversion from a `ParameterBeam` to a `ParticleBeam` (see #408) (@jank324, @Hespe)
 
 ### 🐆 Other

cheetah/accelerator/segment.py

@@ -304,7 +304,7 @@ def from_ocelot(
         :param cell: Ocelot cell, i.e. a list of Ocelot elements to be converted.
         :param name: Unique identifier for the entire segment.
         :param warnings: Whether to print warnings when objects are not supported by
-            Cheetah or converted with potentially unexpected behavior.
+            Cheetah or converted with potentially unexpected behaviour.
         :param device: Device to place the lattice elements on.
         :param dtype: Data type to use for the lattice elements.
         :return: Cheetah segment closely resembling the Ocelot cell.
@@ -324,6 +324,7 @@ def from_bmad(
         cls,
         bmad_lattice_file_path: str,
         environment_variables: dict | None = None,
+        warnings: bool = True,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> "Segment":
@@ -338,13 +339,15 @@ def from_bmad(
         :param bmad_lattice_file_path: Path to the Bmad lattice file.
         :param environment_variables: Dictionary of environment variables to use when
             parsing the lattice file.
+        :param warnings: Whether to print warnings when elements or expressions are not
+            supported by Cheetah or converted with potentially unexpected behaviour.
         :param device: Device to place the lattice elements on.
         :param dtype: Data type to use for the lattice elements.
         :return: Cheetah `Segment` representing the Bmad lattice.
         """
         bmad_lattice_file_path = Path(bmad_lattice_file_path)
         return bmad.convert_lattice_to_cheetah(
-            bmad_lattice_file_path, environment_variables, device, dtype
+            bmad_lattice_file_path, environment_variables, warnings, device, dtype
         )
 
     @classmethod

cheetah/converters/bmad.py

@@ -185,9 +185,10 @@ def convert_element(
             return cheetah.Cavity(
                 length=torch.tensor(bmad_parsed["l"], **factory_kwargs),
                 voltage=torch.tensor(bmad_parsed.get("voltage", 0.0), **factory_kwargs),
-                phase=torch.tensor(
-                    -torch.rad2deg(bmad_parsed.get("phi0", 0.0) * 2 * torch.pi),
-                    **factory_kwargs,
+                phase=-torch.rad2deg(
+                    torch.tensor(bmad_parsed.get("phi0", 0.0), **factory_kwargs)
+                    * 2
+                    * torch.pi
                 ),
                 frequency=torch.tensor(bmad_parsed["rf_frequency"], **factory_kwargs),
                 name=name,
@@ -285,6 +286,7 @@ def convert_element(
 def convert_lattice_to_cheetah(
     bmad_lattice_file_path: Path,
     environment_variables: dict | None = None,
+    warnings: bool = True,
     device: torch.device | None = None,
     dtype: torch.dtype | None = None,
 ) -> "cheetah.Element":
@@ -299,6 +301,8 @@ def convert_lattice_to_cheetah(
     :param bmad_lattice_file_path: Path to the Bmad lattice file.
     :param environment_variables: Dictionary of environment variables to use when
         parsing the lattice file.
+    :param warnings: Whether to print warnings when elements or expressions are not
+            supported by Cheetah or converted with potentially unexpected behaviour.
     :param device: Device to use for the lattice. If `None`, the current default device
         of PyTorch is used.
     :param dtype: Data type to use for the lattice. If `None`, the current default dtype
@@ -337,7 +341,7 @@ def convert_lattice_to_cheetah(
     ), "Merging lines should never produce more lines than there were before."
 
     # Parse the lattice file(s), i.e. basically execute them
-    context = parse_lines(merged_lines)
+    context = parse_lines(merged_lines, warnings)
 
     # Convert the parsed lattice info to Cheetah elements
     return convert_element(context["__use__"], context, device, dtype)

cheetah/converters/utils/fortran_namelist.py

@@ -1,3 +1,4 @@
+import math
 import os
 import re
 from copy import deepcopy
@@ -94,13 +95,15 @@ def merge_delimiter_continued_lines(
     return merged_lines
 
 
-def evaluate_expression(expression: str, context: dict) -> Any:
+def evaluate_expression(expression: str, context: dict, warnings: bool = True) -> Any:
     """
     Evaluate an expression in the context of a dictionary of variables.
 
     :param expression: Expression to evaluate.
     :param context: Dictionary of variables to evaluate the expression in the context
         of.
+    :param warnings: Whether to print warnings for unrecognised expressions that might
+        lead to unexpected behaviour when parsed as strings.
     :return: Result of evaluating the expression.
     """
 
@@ -146,10 +149,11 @@ def evaluate_expression(expression: str, context: dict) -> Any:
 
             return eval(expression, context)
         except SyntaxError:
-            print(
-                f"WARNING: Could not evaluate expression {expression}. It will now be "
-                f"treated as a string. This may lead to unexpected behaviour."
-            )
+            if warnings:
+                print(
+                    f"WARNING: Could not evaluate expression {expression}. It will now "
+                    "be treated as a string. This may lead to unexpected behaviour."
+                )
             return expression
         except Exception as e:
             print(expression)
@@ -180,13 +184,15 @@ def resolve_object_name_wildcard(wildcard_pattern: str, context: dict) -> list:
     return type_matching_keys
 
 
-def assign_property(line: str, context: dict) -> dict:
+def assign_property(line: str, context: dict, warnings: bool = True) -> dict:
     """
     Assign a property of an element to the context.
 
     :param line: Line of a property assignment to be parsed.
     :param context: Dictionary of variables to assign the property to and from which to
         read variables.
+    :param warnings: Whether to print warnings for unrecognised expressions that might
+        lead to unexpected behaviour when parsed as strings.
     :return: Updated context.
     """
     pattern = r"([a-z0-9_\*:]+)\[([a-z0-9_%]+)\]\s*=(.*)"
@@ -201,7 +207,7 @@ def assign_property(line: str, context: dict) -> dict:
     else:
         object_names = [object_name]
 
-    expression_result = evaluate_expression(property_expression, context)
+    expression_result = evaluate_expression(property_expression, context, warnings)
 
     for name in object_names:
         if name not in context:
@@ -211,13 +217,15 @@ def assign_property(line: str, context: dict) -> dict:
     return context
 
 
-def assign_variable(line: str, context: dict) -> dict:
+def assign_variable(line: str, context: dict, warnings: bool = True) -> dict:
     """
     Assign a variable to the context.
 
     :param line: Line of a variable assignment to be parsed.
     :param context: Dictionary of variables to assign the variable to and from which to
         read variables.
+    :param warnings: Whether to print warnings for unrecognised expressions that might
+        lead to unexpected behaviour when parsed as strings.
     :return: Updated context.
     """
     pattern = r"([a-z0-9_]+)\s*=(.*)"
@@ -226,18 +234,20 @@ def assign_variable(line: str, context: dict) -> dict:
     variable_name = match.group(1).strip()
     variable_expression = match.group(2).strip()
 
-    context[variable_name] = evaluate_expression(variable_expression, context)
+    context[variable_name] = evaluate_expression(variable_expression, context, warnings)
 
     return context
 
 
-def define_element(line: str, context: dict) -> dict:
+def define_element(line: str, context: dict, warnings: bool = True) -> dict:
     """
     Define an element in the context.
 
     :param line: Line of an element definition to be parsed.
     :param context: Dictionary of variables to define the element in and from which to
         read variables.
+    :param warnings: Whether to print warnings for unrecognised expressions that might
+        lead to unexpected behaviour when parsed as strings.
     :return: Updated context.
     """
     pattern = r"([a-z0-9_\.]+)\s*\:\s*([a-z0-9_]+)(\s*\,(.*))?"
@@ -267,7 +277,7 @@ def define_element(line: str, context: dict) -> dict:
             property_expression = property_expression.strip()
 
             element_properties[property_name] = evaluate_expression(
-                property_expression, context
+                property_expression, context, warnings
             )
 
     context[element_name] = element_properties
@@ -366,12 +376,14 @@ def parse_use_line(line: str, context: dict) -> dict:
     return context
 
 
-def parse_lines(lines: str) -> dict:
+def parse_lines(lines: str, warnings: bool = True) -> dict:
     """
     Parse a list of lines from a Bmad or Elegant lattice file. They should be cleaned
     and merged before being passed to this function.
 
     :param lines: List of lines to parse.
+    :param warnings: Whether to print warnings for unrecognised expressions that might
+        lead to unexpected behaviour when parsed as strings.
     :return: Dictionary of variables defined in the lattice file.
     """
     property_assignment_pattern = r"[a-z0-9_\*:]+\[[a-z0-9_%]+\]\s*=.*"
@@ -389,21 +401,25 @@ def parse_lines(lines: str) -> dict:
         "m_electron": (
             physical_constants["electron mass energy equivalent in MeV"][0] * 1e6
         ),
-        "raddeg": scipy.constants.degree,
+        "sqrt": math.sqrt,
+        "asin": math.asin,
+        "sin": math.sin,
+        "cos": math.cos,
         "abs_func": abs,
+        "raddeg": scipy.constants.degree,
     }
 
     for line in lines:
         if re.fullmatch(property_assignment_pattern, line):
-            context = assign_property(line, context)
+            context = assign_property(line, context, warnings)
         elif re.fullmatch(variable_assignment_pattern, line):
-            context = assign_variable(line, context)
+            context = assign_variable(line, context, warnings)
         elif re.fullmatch(line_definition_pattern, line):
             context = define_line(line, context)
         elif re.fullmatch(overlay_definition_pattern, line):
             context = define_overlay(line, context)
         elif re.fullmatch(element_definition_pattern, line):
-            context = define_element(line, context)
+            context = define_element(line, context, warnings)
         elif re.fullmatch(use_line_pattern, line):
             context = parse_use_line(line, context)
 

cheetah/converters/utils/rpn.py

@@ -12,7 +12,7 @@ def evaluate_expression(expression: str, context: dict | None = None) -> Any:
 
     stack = []
     stripped = expression.strip()
-    for token in stripped.split(" "):
+    for token in stripped.split():
         match token:
             case "+":
                 try:
@@ -99,12 +99,12 @@ def evaluate_expression(expression: str, context: dict | None = None) -> Any:
                     raise SyntaxError(
                         f"Invalid expression: {expression} - Need one value before asin"
                     )
-            case _:  # all other tokens
-                # commment, ignore this and all following tokens
+            case _:  # All other tokens
+                # Commment ... ignore this and all following tokens
                 if token[0] == "#":
                     break
                 try:
-                    # read as float since it's all torch in the back anyway
+                    # Read as float since it's all torch in the back anyway
                     number = float(token)
                 except ValueError:
                     if token in context:

tests/resources/bmad_tutorial_lattice.bmad

@@ -10,7 +10,9 @@ parameter[particle] = electron  ! Reference particle.
 abs = -0.6
 
 d: drift, L = 0.5 * (0.3 + 0.7)
-b: sbend, L = 0.5, g = 1, e1 = 0.1, dg = 0.001 ! g = 1/design_radius
+! The two spaces in the expression for L are on purpose to test the parser's ability to
+! handle this
+b: sbend, L = 0.6  -0.1, g = 1, e1 = 0.1, dg = sqrt(0.000001) ! g = 1/design_radius
 q: quadrupole, L = abs(abs), k1 = 0.23
 
 lat: line = (d, b, q) ! List of lattice elements