Merge pull request #12 from JPTS2/0.1.9

0.1.9

Author: JPTS2

sad2xs/config.py

@@ -148,7 +148,8 @@ class Config:
         default_factory = lambda: {
         "Marker", "Drift",
         "Bend", "Quadrupole", "Sextupole", "Octupole", "Multipole", "Solenoid",
-        "Cavity", "XYShift", "ZetaShift", "XRotation", "YRotation", "SRotation"})
+        "Cavity", "XYShift", "ZetaShift", "XRotation", "YRotation", "SRotation",
+        "LimitEllipse", "LimitRect"})
 
     ########################################
     # Marker Insertion Tolerance

sad2xs/converter/_004_element_converter.py

@@ -3,7 +3,7 @@
 =============================================
 Author(s):  John P T Salvesen
 Email:      [email protected]
-Date:       09-10-2025
+Date:       18-11-2025
 """
 
 ################################################################################
@@ -1300,25 +1300,132 @@ def convert_apertures(parsed_elements, environment):
         ########################################
         # Initialise parameters
         ########################################
-        a       = 1.0
-        b       = 1.0
+        offset_x    = 0.0
+        offset_y    = 0.0
+        a           = None
+        b           = None
+        dx1         = None
+        dx2         = None
+        dy1         = None
+        dy2         = None
+        aper_type   = None
 
         ########################################
         # Read values
         ########################################
+        if 'dx' in ele_vars:
+            offset_x    = parse_expression(ele_vars['dx'])
+        if 'dy' in ele_vars:
+            offset_y    = parse_expression(ele_vars['dy'])
         if 'ax' in ele_vars:
             a = parse_expression(ele_vars['ax'])
         if 'ay' in ele_vars:
             b = parse_expression(ele_vars['ay'])
+        if "dx1" in ele_vars:
+            dx1 = parse_expression(ele_vars['dx1'])
+        if "dx2" in ele_vars:
+            dx2 = parse_expression(ele_vars['dx2'])
+        if "dy1" in ele_vars:
+            dy1 = parse_expression(ele_vars['dy1'])
+        if "dy2" in ele_vars:
+            dy2 = parse_expression(ele_vars['dy2'])
+
+        ########################################
+        # Determine type of aperture
+        ########################################
+        if any(v is not None for v in [dx1, dx2, dy1, dy2]) and \
+                any(v is not None for v in [a, b]):
+            raise ValueError(
+                f"Error! Aperture {ele_name} has both rectangular and elliptical definitions. This is not supported.")
+        elif any(v is not None for v in [dx1, dx2, dy1, dy2]):
+            aper_type   = 'LimitRect'
+            
+            if dx1 is None and dx2 is None:
+                dx1 = -1.0
+                dx2 = +1.0
+            elif dx1 is None and isinstance(dx2, float):
+                if dx2 < 0:
+                    dx1 = dx2
+                    dx2 = +1.0
+                else:
+                    dx1 = -1.0
+            elif dx2 is None and isinstance(dx1, float):
+                if dx1 < 0:
+                    dx2 = +1.0
+                else:
+                    dx2 = dx1
+                    dx1 = -1.0
+            elif isinstance(dx1, float) and isinstance(dx2, float):
+                if dx1 > dx2:
+                    tmp = dx1
+                    dx1 = dx2
+                    dx2 = tmp
+            else:
+                # At least one is expression, cannot compare
+                # This might cause issues
+                pass
+
+            if dy1 is None and dy2 is None:
+                dy1 = -1.0
+                dy2 = +1.0
+            elif dy1 is None and isinstance(dy2, float):
+                if dy2 < 0:
+                    dy1 = dy2
+                    dy2 = +1.0
+                else:
+                    dy1 = -1.0
+            elif dy2 is None and isinstance(dy1, float):
+                if dy1 < 0:
+                    dy2 = +1.0
+                else:
+                    dy2 = dy1
+                    dy1 = -1.0
+            elif isinstance(dy1, float) and isinstance(dy2, float):
+                if dy1 > dy2:
+                    tmp = dy1
+                    dy1 = dy2
+                    dy2 = tmp
+            else:
+                # At least one is expression, cannot compare
+                # This might cause issues
+                pass
+
+            
+        elif any(v is not None for v in [a, b]):
+            aper_type   = 'LimitEllipse'
+
+            if a is None:
+                a = 1.0
+            if b is None:
+                b = 1.0
+
+        else:
+            raise ValueError(f"Error! Aperture {ele_name} has no valid definition.")
 
         ########################################
         # Create Element
         ########################################
-        environment.new(
-            name    = ele_name,
-            parent  = xt.LimitEllipse,
-            a       = a,
-            b       = b)
+        if aper_type == 'LimitRect':
+            environment.new(
+                name    = ele_name,
+                parent  = xt.LimitRect,
+                min_x   = dx1,
+                max_x   = dx2,
+                min_y   = dy1,
+                max_y   = dy2,
+                shift_x = offset_x,
+                shift_y = offset_y)
+            continue
+        elif aper_type == 'LimitEllipse':
+            environment.new(
+                name    = ele_name,
+                parent  = xt.LimitEllipse,
+                a       = a,
+                b       = b,
+                shift_x = offset_x,
+                shift_y = offset_y)
+        else:
+            raise ValueError(f"Error! Aperture {ele_name} has unsupported definition.")
         continue
 
 ################################################################################

sad2xs/converter/_005_line_converter.py

@@ -162,15 +162,24 @@ def convert_lines(
             # Line and not from the importer: generated line
             # This is done to handle solenoids, ref shifts, thick cavities etc
             if '-' in component \
-                    and component[1:] not in parsed_lines\
+                    and component[1:] not in parsed_lines \
                     and component[1:] in environment.lines:
-                
+                # Checks for:
+                #   - negative sign
+                #   - The line is generated, not imported (parsed lines)
+                #   - The line exists in the environment (to be reversed)
+
                 reversed_line_name      = component[1:] + '_reversed'
+                
+                # Check if the line hasn't already been reversed (duplicate element)
+                if reversed_line_name in environment.lines:
+                    components[i] = reversed_line_name
+                    continue
+                
                 reversed_line_elements  = environment.lines[component[1:]].element_names
 
                 # If it is a generated subline, do not reverse the order of the elements
-
-                # Negate the individual elements
+                # Just negate the individual elements
                 reversed_line_elements  = [f'-{elem}' for elem in reversed_line_elements]
 
                 reverse_handled_components  = []

sad2xs/converter/_010_write_lattice.py

@@ -24,10 +24,11 @@
 from ..output_writer._008_sol import create_solenoid_lattice_file_information
 from ..output_writer._009_cavity import create_cavity_lattice_file_information
 from ..output_writer._010_refshift import create_refshift_lattice_file_information
-from ..output_writer._011_marker import create_marker_lattice_file_information
-from ..output_writer._012_line import create_line_lattice_file_information
-from ..output_writer._013_model import create_model_lattice_file_information
-from ..output_writer._014_offset_markers import create_offset_marker_lattice_file_information
+from ..output_writer._011_aperture import create_aperture_lattice_file_information
+from ..output_writer._012_marker import create_marker_lattice_file_information
+from ..output_writer._013_line import create_line_lattice_file_information
+from ..output_writer._014_model import create_model_lattice_file_information
+from ..output_writer._015_offset_markers import create_offset_marker_lattice_file_information
 
 today   = date.today()
 
@@ -131,6 +132,24 @@ def write_lattice(
     ########################################
     line_table  = line.get_table(attr = True)
 
+    ########################################
+    # Prepare for removal of - signs where not needed
+    ########################################
+    element_names   = line_table.name
+
+    minus_elements  = line_table.rows["-.*"].name
+    for minus_element in minus_elements:
+        root_name   = minus_element.split("::")[0][1:]
+        plus_eles   = [name.startswith(root_name) for name in element_names]
+
+        if any(plus_eles):
+            plus_name   = element_names[plus_eles][0]
+            type_minus  = line_table["element_type", minus_element]
+            type_plus   = line_table["element_type", plus_name]
+
+            assert type_minus == type_plus, \
+                f"Element types for element and its negative do not match"
+
     ########################################
     # Drifts
     ########################################
@@ -210,6 +229,14 @@ def write_lattice(
         line        = line,
         line_table  = line_table,
         config      = config)
+    
+    ########################################
+    # Apertures
+    ########################################
+    lattice_file_string += create_aperture_lattice_file_information(
+        line        = line,
+        line_table  = line_table,
+        config      = config)
 
     ########################################
     # Markers

sad2xs/main.py

@@ -40,6 +40,7 @@ def convert_sad_to_xsuite(
         reverse_element_order:          bool        = False,
         reverse_bend_direction:         bool        = False,
         reverse_charge:                 bool        = False,
+        install_apertures_as_markers:   bool        = False,
         **kwargs):
 
     ############################################################################
@@ -74,6 +75,25 @@ def convert_sad_to_xsuite(
         parsed_lattice_data = parsed_lattice_data,
         excluded_elements   = excluded_elements,
         config              = config)
+    
+    ############################################################################
+    # Check if apertures should become markers
+    ############################################################################
+    if install_apertures_as_markers:
+        if config._verbose:
+            print_section_heading("Converting apertures to markers", mode = 'section')
+                    
+            if "apert" in parsed_lattice_data['elements']:
+                if "mark" in parsed_lattice_data['elements']:
+                    merged = {
+                        **parsed_lattice_data['elements']["apert"],
+                        **parsed_lattice_data['elements']["mark"]}    # Mark takes precedence
+                    parsed_lattice_data['elements']["mark"] = merged
+                else:
+                    parsed_lattice_data['elements']["mark"] = \
+                        parsed_lattice_data['elements']["apert"]
+                
+                parsed_lattice_data['elements'].pop("apert")
 
     ############################################################################
     # Build Environment

sad2xs/output_writer/_002_bend.py

@@ -27,7 +27,8 @@ def create_bend_lattice_file_information(
     ########################################
     # Get information
     ########################################
-    hbends, vbends, sbends, unique_bend_variables, bend_name_dict = extract_bend_information(line, line_table)
+    hbends, vbends, sbends, unique_bend_variables, bend_name_dict = \
+        extract_bend_information(line, line_table)
 
     hbend_lengths       = np.array(sorted(hbends.keys()))
     hbend_names         = generate_magnet_for_replication_names(hbends, "hbend")
@@ -85,6 +86,12 @@ def create_bend_lattice_file_information(
         for replica_name in hbends[hbend_length]:
             replica_variable    = bend_name_dict[replica_name]
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in hbends[hbend_length]:
+                    replica_name        = root_name
+
             # If simple try to make it more compact
             if check_is_simple_bend_corr(line, replica_name):
                 bend_generation = f"""
@@ -129,6 +136,12 @@ def create_bend_lattice_file_information(
         for replica_name in vbends[vbend_length]:
             replica_variable    = bend_name_dict[replica_name]
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in vbends[vbend_length]:
+                    replica_name        = root_name
+
             # If simple try to make it more compact
             if check_is_simple_bend_corr(line, replica_name):
                 bend_generation = f"""
@@ -172,6 +185,12 @@ def create_bend_lattice_file_information(
         for replica_name in sbends[sbend_length]:
             replica_variable    = bend_name_dict[replica_name]
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in sbends[sbend_length]:
+                    replica_name        = root_name
+
             # If simple try to make it more compact
             if check_is_simple_bend_corr(line, replica_name):
                 bend_generation = f"""

sad2xs/output_writer/_003_corr.py

@@ -86,6 +86,12 @@ def create_corrector_lattice_file_information(
         for replica_name in hcorrs[hcorr_length]:
             replica_variable    = corr_name_dict[replica_name]
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in hcorrs[hcorr_length]:
+                    replica_name        = root_name
+
             # If simple try to make it more compact
             if check_is_simple_bend_corr(line, replica_name):
                 corr_generation = f"""
@@ -129,6 +135,12 @@ def create_corrector_lattice_file_information(
         for replica_name in vcorrs[vcorr_length]:
             replica_variable    = corr_name_dict[replica_name]
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in vcorrs[vcorr_length]:
+                    replica_name        = root_name
+
             # If simple try to make it more compact
             if check_is_simple_bend_corr(line, replica_name):
                 corr_generation = f"""
@@ -171,6 +183,12 @@ def create_corrector_lattice_file_information(
         for replica_name in scorrs[scorr_length]:
             replica_variable    = corr_name_dict[replica_name]
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in scorrs[scorr_length]:
+                    replica_name        = root_name
+
             # If simple try to make it more compact
             if check_is_simple_bend_corr(line, replica_name):
                 corr_generation = f"""

sad2xs/output_writer/_004_quad.py

@@ -75,6 +75,12 @@ def create_quadrupole_lattice_file_information(
     for quad, quad_length in zip(quad_names, quad_lengths):
         for replica_name in quads[quad_length]:
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in quads[quad_length]:
+                    replica_name        = root_name
+
             if check_is_simple_quad_sext_oct(line, replica_name, "Quadrupole"):
 
                 if not check_is_skew_quad_sext_oct(line, replica_name, "Quadrupole"):

sad2xs/output_writer/_005_sext.py

@@ -75,6 +75,12 @@ def create_sextupole_lattice_file_information(
     for sext, sext_length in zip(sext_names, sext_lengths):
         for replica_name in sexts[sext_length]:
 
+            # Remove the minus sign if no non minus version exists
+            if replica_name.startswith("-"):
+                root_name   = replica_name[1:]
+                if root_name not in sexts[sext_length]:
+                    replica_name        = root_name
+
             if check_is_simple_quad_sext_oct(line, replica_name, "Sextupole"):
 
                 if not check_is_skew_quad_sext_oct(line, replica_name, "Sextupole"):