Create StagedModel class for chaining LUMEModel instances

[roussel-ryan]

This pull request introduces a new StagedModel class in lume/staged_model.py that enables the composition of multiple LUMEModel instances into a staged sequence. This allows users to chain together several models, passing beam distributions between them in a controlled and validated manner.

The most important changes are:

New Feature: Staged Model Composition

• Added the StagedModel class, which allows users to stage multiple LUMEModel instances in sequence, automatically passing beam distributions between them. (lume/staged_model.py)

Validation and Safety

• Implemented a validate_lume_model_instances class method to ensure that all staged models have compatible beam_input and beam_output variables of type ParticleGroupVariable, preventing misconfiguration. (lume/staged_model.py)

API and Interface

• Provided a supported_variables property that aggregates the supported variables from all staged models, with a note about potential naming conflicts. (lume/staged_model.py)
• Overrode the _set method to propagate input values through the staged models, updating and passing beam distributions as needed. (lume/staged_model.py)

[electronsandstuff]

Overall looks good.

• For the input and output beam variables. I am a little worried about having them hard-coded. This limits flexibility to use with pre-packaged models that have beams exposed, but with the wrong names. Especially with a high level object in lume-base I would prefer to not tie users into a particular naming scheme. Could we look at defining them in software.
• Using fixed names also has the issue that they now conflict with each other with multiple stages... ie beam_input exists as a variable in all stages except maybe the first. I think this deserves some thought.

[electronsandstuff]

This could go in as a validation class method that gets called in __init__.

[roussel-ryan]

Thanks for the comment @electronsandstuff, I'm happy to hear any ideas, we can also chat tomorrow about it

[electronsandstuff]

OK, some thoughts on how this could be structured. The big worry I have is that the way the code is structured requires non-overlapping sets of variable names in each stage for things to be unambiguous. However, at the same time, it also requires that the stages have the shared names for initial and final particles.

One way this could be resolved is to define the particle input / output programmatically.

class ModelParticleInputMixIn:
  @property
  def input_particles(self):
    ...

  @input_particles.setter
  def input_particles(self, par):
    ...

class ModelParticleOutputMixIn:
  @property
  def output_particles(self):
    ...

These classes can get added by inheritance to any existing concrete lume-models that have a clear input and output particle variable. For instance, in Tao, there is a clearly defined track_start and track_end. Impact-T, I believe, also has input and output particles defined in lume-impact already. For other cases where there isn't a clear input and output particle group defined in the code and it is more user defined, we can required users to define the properties themselves to point where they want, the concrete objects in this case could also have a field to have the input or output particles defined by users that way. Users can also override the method in their code to change the behavior.

If this seems like an OK implementation, we would add these to staged model. Then the concrete model classes would have the mixins added and the input and output particles defined if they are clear from the simulation tool. The staged model class would validate that each class has the required interface (ie that the first is an instance of ModelParticleOutputMixIn, the last is of ModelParticleInputMixIn and the rest have both). We would also validate that the models have non-overlapping variable names or have some conflict avoidance mechanism.