Modular API

[benjaminpope]

Hi all,

Quick thread musing on possible modular API patterns.

I would like the information about the data to live in a single equinox OIData module, which will be aware of the many possibilities for interferometric data, including possibilities for

• {visibilities, square visibilities}
• {absolute phases, closure phases, no phases}
• wavelength channels

Likewise, we might fit a Source combination of

• point sources
• resolved sources (Gaussians, annuli),
• more complex sources implemented elsewhere in Jax, eg
  • Harmonix sources,
  • dLux sources,
  • oblate stars....

As a suggestion, the SourceModel should be a equinox module that evaluates the complex visibilities at arguments (u, v, wavel). You should be able to assemble a composite source out of a list of these and it will be appropriately normalized. The OIData module should have a method to take complex visibilities and put them in the same form as the data.

The likelihood function then can be either a method of OIData, or use some nicer syntax in numpyro maybe, which can be abstract and rely on the parametrization in the Source and OIData modules for its details. Having a closed form likelihood function is good for the vmap and Laplace syntax we want to do to implement fast CANDID-like models, which might argue for an OIData method at least for Gaussian uncertainties. We would aim to do sampling with blackjax then, but potentially numpyro.

What do people think? Is there a trivial way to do this in Zodiax?

[benjaminpope]

Consider a likelihood syntax like

oidata = OIData(fname_oifits) # recognizes what the visibility and phase information format is

model = BinaryModel(star_flux, position) + DustModel(dust_flux) # models add with flux scaling and phase offset for position to be another Model instance

def log_prob(params,data,model):
    # some Zodiax syntax here to set the model params 
    model.set(path,*params) 

    # evaluate model complex visibilities at u, v, wavel 
    # then do the likelihood wrt the actual data appropriately
    return oidata.likelihood(model) 

closure_phase etc are now methods of OIData set dynamically at init time, and Model instances have an __add__ method or similar that adds flux & position scaled models to return a new Model.

[benjaminpope]

A minimal working example of this has been implemented in PR #11.