General interface for Background & Perturbation Emulators?

[arthurmloureiro]

Hi @cloe-org/cloe-maintainers !

I have been discussing with @ivansladoljev, @alessiospuriomancini and @marcobonici if it is possible to create a generic interface/api (maybe a protocol?) for emulators in cloelib.

This would avoid us having to code new interfaces for each new emulator that we want to add to cloelib. @ivansladoljev 's PR (#255) is a step towards that but I still think we could bash our heads together to find a solution that is agnostic of:

• Architecture
• Underlying Model (and number of parameters)
• Trained weights
• Parameter ranges
• Overall framework that the emulator has been trained (jax, torch, tensorflow, etc).

I do not have enough experience with creating emulators to know if this is even possible, but I imagine so.

We could have a generic interface that reads a yaml file containing paths to weights/architectures and dictionary of ranges and the same API to return the backgrounds and perturbations. Does that make sense?

Something like:

MySuperDuperEmulator:
    emulator_architecture: "path/to/architecture.pkl"
    emulator_weights: "path/to/trained_weights.pkl"
    returns:
        - pk_linear
        - pk_nonlinear
        - pk_cb_linear
        - pk_cb_nonlinear
        - growth
    ranges:
        - omega_c: [0.01, 0.5]
        - omega_b: [0.001, 0.2]
        - w0: [-3, 3]
        - wa: [...]
        - etc

then the interface is instantiated with something like (just a brainstorming):

from cloelib.cosmology.emulators import (
	EmulatorAPI, 
	EmulatorBackground, 
	EmulatorLinearPerturbations, 
	EmulatorNonLinearPerturbations
)

my_emulator = EmulatorAPI("path/to/my_super_duper_emulator.yaml")

emulator_background = EmulatorBackground(
	[... cosmology goes here ...]
	emulator = my_emulator # so we can test if cosmology within given range
)

emulator_linear = EmulatorLinearPerturbations(
	background = emulator_background,
	redshifts = zs,
	emulator = my_emulator,
)

emulator_nonlinear = EmulatorNonLinearPerturbations(
	background = emulator_background,
	linearperturbations= emulator_linear,
	redshifts = zs,
	emulator = my_emulator,
)

This way we can generically support any emulator with a single interface instead of coding new interfaces for each new emulator.

[gcanasherrera]

While reviewing @ivansladoljev’s PR, I’ve been giving this some thought. I think the idea is excellent, but it may extend beyond the intended scope of cloelib itself. The challenge is that these emulators are not the only ones we foresee interfacing with cloelib.

I’m wondering whether this Protocol/Interface/Class structure might make more sense as a standalone component — effectively a dedicated emulator API — while keeping the cloelib Perturbation protocol as the common interface for all possible codes.

I’d be very happy to support @ivansladoljev in making this happen!

[arthurmloureiro]

I've been thinking about that that being some standalone thing too...

[ivansladoljev]

I am happy to be a part of this work ! I feel that if we put a little effort into this now, it will save us all a lot of stress later on.

[gcanasherrera]

Don't be shy, just open a repo with a catchy name in pypi that is not taken and go for it, happy to help where I can!

[arthurmloureiro]

@gcanasherrera Should this be done inside the umbrella of cloe-org?

The idea is to avoid having to change a thousand cloelib/cosmology/my_cute_new_emulator_N+\infty.py if anything is changed in the cosmology protocol, for example as well as avoiding new my_cute_new_emulator_N+\infty.py PRs because it will for sure get out of hand quite quickly (just think of how many emulators we are talking about today at JC1 telecon 😖 )