Refactor input transform attachment

[swinter1]

Background

• Ax transform the problem from the input/problem space, to the model space (unit cube x, standardized y etc.).
• The GPs built by Botorch are in the model space (the see only standardised inputs)
• The GP need to calcualted QoIs which are in the problem space, so the GP needs to get information about how to:
  • transform x data (passed by the dataloader) to the model space
  • transform y results from the GP into the problem space.

Issue:

• The current solution is clunky, and requiring us to create a model_bridge, find the transforms applied, convert these to botorch transforms, and attach them to the QoI.
• There is also a risk that the transforms on the QoI get out of sync with the model (resulting in using transforms being used that aren't applicable for the model). While the risk is low-ish, it will be very hard to spot this bug.

Business need.

This makes our code messy and is a required step for the QoIEstimator that is completely undocumented. removing this make our code more maintainable.

Definition of done:

• Explore the different possible solutions for this. implement one that work, or document why they don't and why the current solution has to be kept (so future developer don't waste time trying the same things)

Notes:

• Botorch transforms themselves e.g botorch.models.transforms.outcome use .eval()/.train() to internally control if they fit to data. We should double check this is being handled safely/correctly, as this could be a hard issue to spot otherwise.

possible ideas

Attach the transforms to the GP rather than the QoI

When the GP is passed make sure it has the required transforms attached to it (then QoIEstimator should use the call which automatically applies these - I forget exactly which one this is but it is documented in the code)

Risks

• Double applying transformed (or not applying them at all) causes silent errors - we need to be very careful reviewing this change. Lots of places assume that the GP has no transforms internally.
• how and where is the GP created? Is is possible to attach transforms at this point etc?
• When ax uses this gp internally does it assume it should be in the model space (and that the GP has no transforms), will other ax functionality break because of this change?

Remove all transform application by Ax and shift it to the botorch model.

Can we just make the botorch GP have all the transforms, that is the thing we use everywhere, would make it convient

Risks

• breaks the ax interface (maybe?)

[ClaasRostock]

Hi @swinter1 ,

this is highly interesting. We were also discussing (parts) of this issue the other day in the SimTech team.

Non respective what you do in axtreme, in any case it will be great to stay in touch / exchange.
Hence looping in @Jorgelmh , @davidhjp , @StephanieKemna and @MelihAkdag here as well.

FYI: Some initial ideas and understandings around handling transformations between User Space (Problem Space) and Model Space are also in the mvx Miro board here. Feel free to amend any thoughts and insights also in there.

Best regards

Claas