Skip to content

Add Set-Like Operations to ElementArray #194

@gupichon

Description

@gupichon

##Description, motivation and use case

Description

Introduce native set-like operations on ElementArray (and derived array classes) using Python operators:

  • & → intersection
  • | → union (unique, stable order)
  • + → alias of union
  • - → difference

Additionally, implement automatic result typing based on the resulting element types.


Motivationand use case

In pyAML, array selections are frequently combined to build meaningful subsets of the machine.

For example:

cell1 = sr.live.get_elements("C01")        # ElementArray (mixed types)
sextupoles = sr.live.get_magnets("SEXT")   # MagnetArray

cell1_sext = cell1 & sextupoles

Here:

  • cell1 is an ElementArray containing various element types.
  • sextupoles is a MagnetArray.
  • The intersection keeps only sextupole magnets located in cell 1.
  • The result should be automatically typed as a MagnetArray.

Currently, this requires manual filtering logic. Native operator support would make this natural and expressive.


Proposed Behavior

Supported Operators

Operator Meaning Description
a & b Intersection Elements common to both arrays
a | b Union Unique elements, stable order
a + b Union Alias of |
a - b Difference Elements in a but not in b

Equality Semantics

Elements should be compared based on their name (element.get_name()), not instance identity.

This ensures predictable behavior even if elements originate from different instantiations.


Auto-Typing of Result

The resulting array should be automatically typed:

  • If all remaining elements share the same concrete type, return the corresponding specialized array (e.g., MagnetArray, BPMArray).
  • If elements are mixed types, return a generic ElementArray.
  • If the result is empty, return an empty list (preserving current internal behavior).

Example:

cell1 = sr.live.get_elements("C01")        # ElementArray
sextupoles = sr.live.get_magnets("SEXT")   # MagnetArray

cell1_sext = cell1 & sextupoles
type(cell1_sext).__name__  # -> "MagnetArray"

Peer Consistency

Operations should raise an exception if arrays are attached to different peers (e.g., design vs live).

Example:

design_cell1 & live_cell1  # should raise

Benefits

  • Makes selection composition expressive and natural
  • Reduces boilerplate filtering code
  • Encourages a declarative style for machine subsets
  • Improves usability for both scripting and notebook workflows

Implementation Notes

  • Add operator overloads in ElementArray

  • Reuse existing internal array factory mechanism for result typing

  • Ensure stable order in union

  • Add unit tests covering:

    • Intersection
    • Union
    • Difference
    • Mixed types fallback
    • Peer mismatch error

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions