Skip to content

Ax integrations with the Bluesky ecosystem#112

Merged
thomaswmorris merged 11 commits intobluesky:mainfrom
thopkins32:ax-integration
Jun 19, 2025
Merged

Ax integrations with the Bluesky ecosystem#112
thomaswmorris merged 11 commits intobluesky:mainfrom
thopkins32:ax-integration

Conversation

@thopkins32
Copy link
Copy Markdown
Collaborator

@thopkins32 thopkins32 commented Apr 16, 2025

This PR enables using https://ax.dev/ with the Bluesky ecosystem.

  • Creates some simple functions for using Ax with the Bluesky ecosystem.
  • Adds a unit test using the simulated beamline.
  • Adds a demo notebook showing that optimization works on the simulated KB mirror setup (see the types of visualizations Ax provides).

Closes #93

@review-notebook-app
Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@thopkins32 thopkins32 changed the title Ax Integrations with the Bluesky ecosystem Ax integrations with the Bluesky ecosystem Apr 16, 2025
@thopkins32 thopkins32 marked this pull request as ready for review April 18, 2025 18:49
@thopkins32 thopkins32 requested review from jennmald and mrakitin April 18, 2025 18:49
"metadata": {},
"outputs": [],
"source": [
"for _ in range(25):\n",
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question about Ax, you picked 0-25 as the range for optimization, does more trials improve the optimization process or is it just like an ensemble and averages out at the end?
I see in the markdown you mention the first few are random and rest is BO.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ax actually infers what optimization setup to run based on the parameters you pass in according to the docs: https://ax.readthedocs.io/en/stable/service.html#module-ax.service.ax_client.

If you want something more custom then you can pass in your own GenerationStrategy (which suggests new parameters, or Arms, to try).

"cell_type": "markdown",
"metadata": {},
"source": [
"### Viewing each arm's objective values"
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am a huge fan of the visuals Ax provides.

m = fim > threshold * fim.max()
m = fim > (threshold * fim.max())
area = m.sum()
if area == 0.0:
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why is this needed?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When I had the simulated beamline without noise and the beam was off-screen, the total sum of the image that is above the threshold of 0.5 was 0.0. This leads to a divide by zero error a few lines down. Since the image is all the same value, there is no beam to extract and we can return 0.0 for everything.

@thopkins32 thopkins32 requested a review from jennmald April 22, 2025 17:15
Copy link
Copy Markdown
Contributor

@thomaswmorris thomaswmorris left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

@thomaswmorris thomaswmorris merged commit 59ac2ab into bluesky:main Jun 19, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Ax support

3 participants