Add register() functions for strategies and surrogates by jduerholt · Pull Request #736 · experimental-design/bofire

jduerholt · 2026-02-23T16:39:40Z

Motivation

Provide public API to register custom strategy and surrogate mappings instead of requiring users to mutate internal dictionaries directly. For BotorchSurrogate subclasses, registration automatically rebuilds the AnyBotorchSurrogate union and BotorchSurrogates Pydantic model so custom surrogates work in botorch-based strategies out of the box. It addresses the discussion in this PR: #731

@TobyBoyne @bertiqwerty @LukasHebing @KislayaRavi what do you think?

This is a draft which should serve for discsussion, code build with Claude Code.

Have you read the Contributing Guidelines on pull requests?

Yes.

Have you updated `CHANGELOG.md`?

Not yet.

Test Plan

Unit tests.

Provide public API to register custom strategy and surrogate mappings instead of requiring users to mutate internal dictionaries directly. For BotorchSurrogate subclasses, registration automatically rebuilds the AnyBotorchSurrogate union and BotorchSurrogates Pydantic model so custom surrogates work in botorch-based strategies out of the box. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

TobyBoyne · 2026-02-23T17:24:18Z

I think this looks correct to me! A few minor thoughts:

I would quite like a decorator syntax. So I can write something like below: that way, I can easily see that the new surrogate is registered right next to where it is defined.

@surrogates_api.register(CustomBotorchSurrogateDataModel)
class CustomBotorchSurrogate(Surrogate):
    def __init__(self, data_model, **kwargs):
        ...

Could we initialize SURROGATE_MAP to be empty, and register all of the BoFire models using this new API? This would mean, for example, we would no longer need to maintain the default values in _BOTORCH_SURROGATE_TYPES.
I think the data_model_transform should be a separate thing to register. For example, if I want to have a minimal installation of BoFire (without installing botorch), I would not be able to register a custom transform. Maybe that is okay, since I think it is only ever needed in combination with mapping to an actual strategy, but I still think something that takes as input+output purely data models should be separate from the mapper onto real strategies. But I don't feel this particularly strongly so feel free to ignore.

I think this PR also breaks using AnyBotorchSurrogate in static typing (I'm getting a few new Pylance errors), but it isn't used as a type hint anywhere in the codebase so not a big deal.

…mic Pydantic union rebuilding Extend the register() pattern to kernels, priors, and engineered features mappers so custom types can be registered with both the mapper dictionaries and Pydantic unions. Registration dynamically patches model annotations and triggers model_rebuild() in the correct cascade order (priors → kernels → aggregation kernels → surrogates → BotorchSurrogates) so custom types pass Pydantic validation when used in surrogate model fields. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

@register

…features Replace the manually maintained KERNEL_MAP, PRIOR_MAP, and AGGREGATE_MAP dictionaries with @register() decorators on each map function. The dicts are now initialized empty and populated at import time via the same register() function that external users call for custom types. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

jduerholt · 2026-02-24T09:30:53Z

@TobyBoyne: I was just pushing it further, I need to do an in-depth review here, but is this what you were looking for?

TobyBoyne · 2026-02-24T10:22:03Z

Yes that syntax looks very nice, and I like the decorator being included where the kernels/priors are defined.

I was also wondering why do we use, for example, AnyBotorchSurrogate instead of just BotorchSurrogate? What is the benefit of the union type instead of just type hinting as the parent type? It would avoid needing all of this model rebuilding.

KislayaRavi · 2026-02-24T11:00:21Z

bofire/data_models/kernels/api.py

Do we also need _rebuild_dependent_strategy similar to rebuild_dependent_models if we want full dynamic Pydantic acceptance of new strategy data-model types?

Yes, correct. Should be fixed now. My Agent was overlooking this.

Yes, correct. Should be fixed now. My Agent was overlooking this. (last week)

Haha. Nice try blaming the agent. :D

KislayaRavi · 2026-02-24T11:04:14Z

I just had one comment on the code.
Additionally, I think we should also a Quarto doc tutorial with an example demonstrating steps to externally register kernel or surrogate.

LukasHebing · 2026-02-24T11:56:53Z

Yes that syntax looks very nice, and I like the decorator being included where the kernels/priors are defined.

I was also wondering why do we use, for example, AnyBotorchSurrogate instead of just BotorchSurrogate? What is the benefit of the union type instead of just type hinting as the parent type? It would avoid needing all of this model rebuilding.

I think this is needed for de-serialisation of objects from dicts / json. Having a pydantic class with a field surrogate: AnyBotorchSurrogate makes the de-serialiser go through all classes in AnyBotorchSurrogate, until the fieldnames match.

…strategy rebuild - Extract _patch_field to shared bofire/data_models/_register_utils.py with patch_field() and append_to_union_field() helpers, reused across kernels, priors, and strategies - Make AnyContinuousKernel and AnyCategoricalKernel list-backed and dynamic; register_kernel auto-detects sub-category from ContinuousKernel / CategoricalKernel base class and updates MixedSingleTaskGPSurrogate - Replace lazy imports with data_models.register_kernel() and data_models.register_prior() in mapper register() functions - Remove meta parameter from strategy register (only ACTUAL_MAP) - Add register_strategy() to data models that rebuilds ActualStrategy union, Step, and StepwiseStrategy so custom strategies pass validation - Add introspection tests that verify _rebuild_dependent_models covers all AnyPrior/AnyPriorConstraint/AnyKernel fields in the codebase Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

jduerholt

Some comments from my side.

bofire/data_models/kernels/api.py

bofire/data_models/priors/api.py

bofire/kernels/mapper.py

bofire/priors/mapper.py

bofire/strategies/mapper.py

jduerholt · 2026-02-24T14:39:15Z

bofire/data_models/kernels/api.py

Yes, correct. Should be fixed now. My Agent was overlooking this.

ConditionalEmbeddingKernel.base_kernel and WedgeKernel.base_kernel had hardcoded inline Union types that were not patched by _rebuild_dependent_models, so newly registered kernel types were not accepted. Both parent and child need explicit patching because Pydantic gives each class its own copy of model_fields. Also enhances the introspection test to detect inline kernel Union fields on Kernel container classes, preventing this class of regression. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Resolves conflict in bofire/surrogates/engineered_features.py by combining the register() decorator pattern from feature/register with main's refactored private functions (_weighted_features, _map_weighted_feature, _map_molecular_weighted_feature) and partial() bindings for sum/mean variants. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

jduerholt · 2026-02-24T15:55:24Z

Ok, I did more stuff here, this is now almost final @bertiqwerty can you have a look? what do you think? I will also add a tutorial etc. But would be nice to get your assesment! And of course also the assesment from the others!

bertiqwerty · 2026-02-25T08:55:34Z

Ok, I did more stuff here, this is now almost final @bertiqwerty can you have a look? what do you think? I will also add a tutorial etc. But would be nice to get your assesment! And of course also the assesment from the others!

Sure. I will look into it.

The register() infrastructure uses Union[tuple(list)] to build unions at runtime. This is valid Python but not statically analysable, so tell ty to ignore the diagnostic. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Documents how to use the register() API for strategies, surrogates, kernels, and priors with practical code examples covering both decorator and direct-call forms. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

The `model_fields` attribute is a Pydantic BaseModel class attribute that ty cannot resolve on the generic `type` parameter. Add inline ignore comments on all 5 access sites. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

ruff format splits long lines, moving the ignore comment away from the `.model_fields` access that triggers the error. Place the comment on the first line of each expression so ty sees it on the correct line. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

KislayaRavi · 2026-03-02T12:54:04Z

I faced an issue while running a SoboStrategy campaign using my custom define surrogate. Following code snippet reproduces this error.

from bofire.data_models.surrogates.surrogate import Surrogate as SurrogateDataModel
from bofire.data_models.surrogates.trainable import TrainableSurrogate as TrainableDataModel
import bofire.surrogates.api as surrogates
import bofire.data_models.surrogates.botorch_surrogates as bs

class MySurrogateDataModel(SurrogateDataModel, TrainableDataModel):
    type: Literal["MySurrogate"] = "MySurrogate"

class MySurrogate(Surrogate, TrainableSurrogate):
    def __init__(self, data_model: MySurrogateDataModel, **kwargs):
        super().__init__(data_model=data_model, **kwargs)

surrogates.register(MySurrogateDataModel, MySurrogate)

print(bs.AnyBotorchSurrogate)  # MySurrogate is NOT in the union

Workaround
Temporarily, I found a workaround. Explicitly call register_botorch_surrogate

import bofire.surrogates.api as surrogates
from bofire.data_models.surrogates.botorch_surrogates import register_botorch_surrogate

surrogates.register(MySurrogateDataModel, MySurrogate)
register_botorch_surrogate(MySurrogateDataModel)

But we should fix it.

Probable culprit

# bofire/surrogates/mapper.py
if issubclass(data_model_cls, data_models.BotorchSurrogate):
    register_botorch_surrogate(data_model_cls)  # ← never reached for plain SurrogateDataModel subclasses

jduerholt · 2026-03-03T07:57:43Z

Hmm, but is it actually a bug? The SoboStrategy only works with botorch based surrogates, so your data model should inherit from it, or?

KislayaRavi · 2026-03-03T09:38:48Z

One can use non Botorch Surrogate models with SoboStrategy, provided the attribute model of the custom defined Surrogate has a function posterior.
The ability of define and register own CustomSurrogate was to provide users ability to include non-standard surrogate. That's why the line if issubclass(data_model_cls, data_models.BotorchSurrogate): may not be necessary.
I was able to run a SoboStrategy with my custom surrogate using register() feature added in this PR. I defined my own posterior function. The register() feature is working. The only thing was using register() feature twice, once for the custom data model and other for the corresponding custom surrogate. I feel that it is redundant.

jduerholt · 2026-03-03T11:25:44Z

Hmm, but why not just inherit in the data model from the BotorchSurrogate? There is no problem in this or?

bertiqwerty

Thanks, Johannes.

bertiqwerty · 2026-03-03T12:59:01Z

bofire/data_models/features/api.py

-from typing import Union
+import typing
+from collections.abc import Sequence
+from typing import List, Type, Union


use list instead of List. Don't import List.

bertiqwerty · 2026-03-03T12:59:05Z

bofire/data_models/features/api.py

+AnyEngineeredFeature = Union[tuple(_ENGINEERED_FEATURE_TYPES)]
+
+
+def register_engineered_feature(data_model_cls: Type[EngineeredFeature]) -> None:


Maybe don't clutter the api.py with implementations? Create a separate file and import the function here?

bertiqwerty · 2026-03-03T12:59:11Z

bofire/data_models/features/api.py

+    _ENGINEERED_FEATURE_TYPES.append(data_model_cls)
+    AnyEngineeredFeature = Union[tuple(_ENGINEERED_FEATURE_TYPES)]
+
+    # Lazy import to avoid circular dependencies


see above. if you had separate file for the implementation, would the import still be circular?

bertiqwerty · 2026-03-03T12:59:24Z

bofire/data_models/features/api.py

+    # Patch the Sequence[Union[...]] annotation on EngineeredFeatures.features
+    old = EngineeredFeatures.model_fields["features"].annotation
+    inner_args = typing.get_args(typing.get_args(old)[0])
+    if data_model_cls not in inner_args:


this peace of code looks fragile. but currently i don't have a better idea. hmm... at least if this breaks at some point only the registration functinality is affected, not the existing modules.

bertiqwerty · 2026-03-03T12:59:57Z

bofire/data_models/priors/api.py

@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Union
+from typing import List, Type, Union


bertiqwerty · 2026-03-03T13:00:02Z

bofire/data_models/priors/api.py

+AnyPriorConstraint = Union[tuple(_PRIOR_CONSTRAINT_TYPES)]
+
+
+def _rebuild_dependent_models() -> None:


I prefer a separate file to keep the api.py as clean as possible.

bertiqwerty · 2026-03-03T13:00:08Z

bofire/data_models/strategies/actual_strategy_type.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import List, Type, Union


bertiqwerty · 2026-03-03T13:00:11Z

bofire/data_models/strategies/actual_strategy_type.py

+ActualStrategy = Union[tuple(_ACTUAL_STRATEGY_TYPES)]
+
+
+def register_strategy(data_model_cls: Type[Strategy]) -> None:


again don't clutter the api.py

bertiqwerty · 2026-03-03T13:00:20Z

bofire/kernels/mapper.py

+    Can be used as a decorator or as a direct function call::
+
+        # Decorator form
+        @register(MyKernelDataModel)


Doesn't adding the decorator to existing functions trigger unnecessary rebuild-model calls on import time? Our import time is already pretty long. Further, I find decorators less transparent. If someone wants to register, they could simply use the direct call form. For the existing stuff, I liked the old and stupid dict. And it is somehow incosistent to the data models where we kept the lists of types and extended them. I would prefer to keep the dicts and extend those.

@bertiqwerty: But the decorator in general would be fine? So, for registering stuff outside of the library? Or do you want to remove the decorator support in general?

You mean as option for people to use to add with new code outside of BoFire? There is no way to remove it, is there? The register function can always be used as decorator or with a simple function call, right?

True, so you would like to remove the decorator use from the inner workings of BoFire, correct?

Yes. And I want the dict back.

Ok, I will add the dict back in the internal mappings.

bertiqwerty · 2026-03-03T13:01:25Z

bofire/data_models/kernels/api.py

Yes, correct. Should be fixed now. My Agent was overlooking this. (last week)

Haha. Nice try blaming the agent. :D

KislayaRavi · 2026-03-03T13:50:17Z

Thanks Johanes, it works.
I was inheriting from SurrogateDataModel and not TrainableBotorchSurrogate.

jduerholt · 2026-03-09T13:33:10Z

Hi @bertiqwerty,

I implemented the changes, that you suggested, can you have a look again?

Best,

Johannes

bertiqwerty · 2026-03-09T13:39:21Z

Hi @bertiqwerty,

I implemented the changes, that you suggested, can you have a look again?

Best,

Johannes

My reviews are usually YOLO. Just kidding. Thanks for the changes. I will have a look.

bertiqwerty

Thanks, Johannes.

LukasHebing · 2026-03-10T08:40:51Z

Thanks a lot @jduerholt,
This functionality was really missing in bofire. I think this is really valuable for new users!

jduerholt · 2026-03-10T14:09:01Z

Thanks @LukasHebing: I will merge it tmr, I think.

Since _register.py modules use lazy imports inside function bodies, they can be safely imported at the top of api.py files. Also import register_strategy directly in strategies/api.py instead of re-exporting through actual_strategy_type.py. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Replace manual Sequence[Union[...]] patching with the shared append_to_union_field utility from _register_utils, consistent with the other _register.py modules. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

jduerholt and others added 3 commits February 24, 2026 09:27

decorator syntax

228a1df

KislayaRavi reviewed Feb 24, 2026

View reviewed changes

jduerholt commented Feb 24, 2026

View reviewed changes

jduerholt and others added 2 commits February 24, 2026 16:05

jduerholt marked this pull request as ready for review February 26, 2026 21:21

jduerholt and others added 3 commits February 27, 2026 09:51

Add user guide for registering custom types

3ed1a02

Documents how to use the register() API for strategies, surrogates, kernels, and priors with practical code examples covering both decorator and direct-call forms. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

bertiqwerty self-requested a review March 3, 2026 12:08

bertiqwerty requested changes Mar 3, 2026

View reviewed changes

implement behrangs comments

8da8684

jduerholt requested a review from bertiqwerty March 9, 2026 13:33

bertiqwerty approved these changes Mar 9, 2026

View reviewed changes

jduerholt and others added 2 commits March 12, 2026 19:29

Use append_to_union_field in features/_register.py

954851c

Replace manual Sequence[Union[...]] patching with the shared append_to_union_field utility from _register_utils, consistent with the other _register.py modules. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

jduerholt merged commit a61514b into main Mar 12, 2026
12 checks passed

jduerholt deleted the feature/register branch March 12, 2026 22:07

		AnyEngineeredFeature = Union[tuple(_ENGINEERED_FEATURE_TYPES)]


		def register_engineered_feature(data_model_cls: Type[EngineeredFeature]) -> None:

		AnyPriorConstraint = Union[tuple(_PRIOR_CONSTRAINT_TYPES)]


		def _rebuild_dependent_models() -> None:

		@@ -1,4 +1,4 @@
		from typing import Union
		from typing import List, Type, Union

		ActualStrategy = Union[tuple(_ACTUAL_STRATEGY_TYPES)]


		def register_strategy(data_model_cls: Type[Strategy]) -> None:

Conversation

jduerholt commented Feb 23, 2026

Motivation

Have you read the Contributing Guidelines on pull requests?

Have you updated CHANGELOG.md?

Test Plan

Uh oh!

TobyBoyne commented Feb 23, 2026

Uh oh!

jduerholt commented Feb 24, 2026

Uh oh!

TobyBoyne commented Feb 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bertiqwerty Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

KislayaRavi commented Feb 24, 2026

Uh oh!

LukasHebing commented Feb 24, 2026

Uh oh!

jduerholt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jduerholt commented Feb 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

bertiqwerty commented Feb 25, 2026

Uh oh!

KislayaRavi commented Mar 2, 2026

Uh oh!

jduerholt commented Mar 3, 2026

Uh oh!

KislayaRavi commented Mar 3, 2026

Uh oh!

jduerholt commented Mar 3, 2026

Uh oh!

bertiqwerty left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bertiqwerty Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bertiqwerty Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Have you updated `CHANGELOG.md`?

bertiqwerty Mar 3, 2026 •

edited

Loading

jduerholt commented Feb 24, 2026 •

edited

Loading

bertiqwerty Mar 3, 2026 •

edited

Loading

bertiqwerty Mar 3, 2026 •

edited

Loading

bertiqwerty Mar 3, 2026 •

edited

Loading