feat(explorer): Implement c-optimal discretisation strategy to reduce gradient variance in error budgeting feature

[adlantz]

🔗 Closed Issues

Closes #283

---

📝 Description

The functionget_error_budget() in deltakit-explorer/src/deltakit_explorer/analysis/error_budget/_budget.py takes a fitting_parameters parameter of type FittingParameters. This parameter has a discretisation_strategy property of type DiscretisationStrategy. In order to calculate the proportionate contribution of each error source to the error budget, we need to calculate the gradient for each error source, this requires sampling a set of points around the defined error rate, fitting the results to a polynomial, then calculating the gradient at that point. The discretisation_strategy property dictates which points should be sampled to calculate that gradient.

Currently we only give the users two strategies, DiscretisationStrategy.LINEAR and DiscretisationStrategy.LOGARITHMIC. These choose a linear and logarithmic spread of points respectively. Neither of these sampling strategies is optimal with respect to the variance of the final gradient calculation.

In this PR we introduce a third strategy DiscretisationStrategy.C_OPTIMAL which uses an optimization algorithm to find the best choice of points to sample in order to minimize the variance of the gradient estimate. This kind of optimization is known as a c-optimal experimental design (https://en.wikipedia.org/wiki/Optimal_experimental_design).

The DiscretisationStrategy.C_OPTIMAL approach consists of two main functions. The first isget_c_optimal_points() which sets up the optimization parameters and runs scipy.optimize.differential_evolution(). The second is _c_optimal_objective(), the objective function passed to the SciPy optimization algorithm. This calculates the variance of the gradient estimation given a set of sampling points.

There are a few things to note.

1. In deltakit-explorer/src/deltakit_explorer/analysis/error_budget/_post_processing.py we update the _compute_logical_error_rate_per_round_from_results(). Instead of grabbing the firsts result from the pandas df query, we sum across all the results. This change is necessary because c-optimal can choose to sample single points multiple times. This is non-breaking for the existing use cases of this function, it still works if only 1 row is returned.
2. The existing discretisation functions get_logarithmic_points() and get_linear_points() expect c, the central point to be a python float. However, currently generate_sweep_parameters() calls them with c as a size 1 numpy column vector. This was not breaking before because get_logarithmic_points() and get_linear_points(). We address this by explicitly flattening central_point in generate_sweep_parameters() before calling a discretisation strategy. We ran unit tests to confirm this didn't break the linear and logarithmic approaches.
3. We hard coded the SciPy optimization algorithm's seed. This allows for a sense of reliability/consistency - a user may not want to see their error budget shifting when running the same code with the same parameters twice. However, it allowing for a random seed can produce better results. It might be worth considering a "random vs deterministic" parameter for the user to toggle.

Here is a run across 10 different seeds demonstrating the improvement of c-optimal over linear and logarithmic. We also considered a d-optimal approach but it was not as effective.

Tests

We mimicked the existing pytest unit test infrastructure and added coverage for the new code in this feature.

---

🚦 Status

In review

---

🛠️ Future Work

This bounds issue #216 is a logical follow up.

---

➕️ Additional Information

This work was done by University of Washington students as part of a course in collaboration with @rolandriver

---

🧾 Release Note

Add c-optimal discretisation strategy to improve error budgeting feature

[CLAassistant]

All committers have signed the CLA.

[nelimee]

@adlantz let me know when I should start reviewing this PR. Ideally, it should be out of draft and with CI passing, but I am fine doing a pre-review if you need one :)

[rolandriver]

Hey @adlantz thank you for this. Yes, you can raise an issue with the description above. In the PR description you can reference the issue by adding "Closes #." and GitHub will create the link.

[adlantz]

@nelimee @rolandriver CI checks are passing - PR is ready for review

[rolandriver]

Thanks @adlantz for this really. This is looking quite nice indeed. Few comments from me and happy to approve soon after they've been addressed.

[rolandriver]

Hey @adlantz would you have time to address review comments and finalise this PR? There had been quite few changes merged on main so please if you could rebase.

[nelimee]

I'll review that PR probably around Thursday, thanks for submitting it!

[rolandriver]

Thank you so much for your contribution @adlantz. This is looking good to me. Approving now but merging upon @nelimee's review and green light.

[adlantz]

@rolandriver No problem, thank you for reviewing. I'll keep an eye on it until it's merged to see if there's any more updates to make.

[nelimee]

Overall I think that this brings a nice new way of choosing points for error-budgeting, but think this implements an approximation of what people usually think about when using "optimal design".

To be more explicit, your method starts by fixing a set of points that are sufficiently close to each-other (_C_OPTIMAL_CANDIDATE_GRID_SIZE which is 50 at the moment) and linearly spaced, and then optimises on which point should be evaluated and how many times it should be evaluated.

I think that this method tends to the original optimal-design idea I had when _C_OPTIMAL_CANDIDATE_GRID_SIZE tends to infinity, i.e., when there is no limit on which point (within the provided interval) can be picked.

That is perfectly fine to me, and implementing the continuous optimisation version will probably be a nice exercise for later.

I'll accept that PR once my comments are addressed.

[adlantz]

@nelimee I addressed your comments, thanks for your review. Let me know if there's anything else you want me to touch up!

[rolandriver]

Hey @adlantz would you have time to implement @nelimee comments? Just to finalise this PR. Thank you 🙏

[adlantz]

@rolandriver @nelimee apologies for being a little slow I was traveling without my laptop for the last 10 days. I will address the final comments and fix the merge conflicts tomorrow!

[adlantz]

@rolandriver @nelimee comments addressed and synced with main. I see a part of the CI failing but I don't think that's an issue I caused.

[nelimee]

@rolandriver @nelimee comments addressed and synced with main. I see a part of the CI failing but I don't think that's an issue I caused.

You are right, that's not due to this PR. I checked, none of the failures are due to this PR. Probably due to an update of pydoclint. We can merge as-is if you are ready and don't see anything else to change in this PR :)

[adlantz]

@rolandriver @nelimee Good to merge! I don't think I have the permissions to do it myself

[rolandriver]

Thank you @adlantz for this contribution!