Design for the recOrder-waveorder interface

[talonchandler]

Seeking comments from @mattersoflight @ziw-liu @ieivanov @edyoshikun.

As I've been refactoring phase reconstructions into phase.py I've been realizing that it's a good time to discuss more details of the design of the recOrder-waveorder interface.

Background and goals: All of our reconstructions take a form similar to "do an expensive precomputation" (e.g. calculate a phase OTF, calculate the background-correction Mueller matrices), then apply that result to the acquired data. Our workflows often require some iteration (did I choose the right settings for my transfer function? what does my transfer function look like? what happens if I apply this transfer function to a different dataset?), so we need workflows that allow quick feedback for prototyping while allowing us to apply a reconstruction to a large number of positions/times/files.

This draft design of a waveorder CLI and its configuration settings/files enables the following core workflows:

1. Calculate a transfer function and save it to a napari-readable .zarr file (for inspection by technical users)
   >> wo calc-tf tf-config.yml -o tf.zarr
   then apply its inverse to a file
   >> wo apply-inv-tf tf.zarr input.zarr -o output.zarr

2. Do both steps in one call with:
   >> wo reconstruct recon-config.yml input.zarr -o output.zarr

I plan to have recOrder 0.4.0 generate the tf-config.yml from its GUI, then make wo tf and wo apply-inv-tf calls to perform the reconstruction. This will allow the user to easily transition from a recOrder-GUI-based workflow to a wo-CLI-based workflow.

I expect recOrder 1.0.0 to be able to read this set of pydantic models (or future variants) to generate much of its GUI.

Main questions:

• Earlier versions of recOrder and it's offline mode faced difficulty in handling the expensive precomputation of transfer functions. This design puts some of the burden for handling transfer functions on the user, but I think the advantages are worth it (everyone can precompute and inspect the TFs) and a high-level wo reconstruct command can provide the same one-step experience if desired. Does this design make sense?
• the pydantic models should provide enough information to auto-populate much of the recOrder GUI...any comments here @ziw-liu?
• I decided to split the complete reconstruct step into calc-tf and apply-inv-tf since I wanted each step to be a verb, and it wasn't obvious what to call the second substep (is the second step reconstruct or are both steps together reconstruct?). I'm welcoming suggestions on my naming choices throughout.

---

Below is a skeleton design for the main pydantic models. More detailed validations/typing will follow after a first iteration here.

##### TRANSFER FUNCTION SETTINGS #####

# Shared by all transfer function calculations
class SharedCalcTFSettings(BaseModel):
    lambda_ill: float
    X_shape: int
    Y_shape: int
    Z_shape: int
    use_gpu: bool
    gpu_id: int


# Settings for birefringence-specific transfer functions
# These are needed for precomputing background-correction Mueller matrices
class BirefringenceCalcTFSettings(BaseModel):
    swing: float
    background_path: str  #
    background_correction_mode: str  #
    pol_scheme = str  # 4/5 state


# Settings for phase-specific transfer functions
# These are needed for computing OTFs
class PhaseCalcTFSettings(BaseModel):
    n_media: float
    NA_illu: float
    NA_det: float
    XY_ps: float
    Z_ps: float
    Z_padding: int


# This is the main configuration file used by `wo calc-tf tf-config.yml -o tf.zarr`
# This file will always have shared_settings, will always have at least one of
# biref_settings or phase_settings, and can have both. Is this class the best
# way to encode this logic? Should I use pydantic to validate that one or the
# other is present?
class CalcTFSettings(BaseModel):
    shared_settings: SharedCalcTFSettings
    biref_settings: BirefringenceCalcTFSettings = None
    phase_settings: PhaseCalcTFSettings = None


##### APPLY INV TRANSFER FUNCTION SETTINGS #####


class BirefringenceApplyInvTFSettings(BaseModel):
    orientation_flip: bool
    orienation_rotate: bool


class PhaseApplyInvTFSettings(BaseModel):
    recon_algorithm: str  # Tikhonov, TV, etc
    recon_params: List[float]  # ??? Each algorithm takes a different number of params,
    # what's the best way to encode this information


# This is the main configuration that goes into `wo apply-inv-tf tf.zarr input.zarr -o output.zarr`
class ApplyInvTFSettings(BaseModel):
    biref_apply_settings = BirefringenceApplyInvTFSettings
    phase_apply_ettings = PhaseApplyInvTFSettings


##### RECONSTRUCT PIPELINE #####


# This is the configuration file used in
# `wo reconstruct recon-config.yml <input-MM-folder> -o output.zarr`
class ReconstructSettings(BaseModel):
    calc_tf_settings = CalcTFSettings
    apply_tf_settings = ApplyInvTFSettings

[ziw-liu]

I plan to have recOrder 0.4.0 generate the tf-config.yml from its GUI, then make wo tf and wo apply-inv-tf calls to perform the reconstruction. This will allow the user to easily transition from a recOrder-GUI-based workflow to a wo-CLI-based workflow.

In this case will recOrder still have a CLI? Or we point users to waveorder from there?

[talonchandler]

In this case will recOrder still have a CLI? Or we point users to waveorder from there?

I'm flexible on where the CLI lives...could be in recOrder or waveorder.

I'm currently leaning towards making it a waveorder CLI and making recOrder's CLI minimal. After each recOrder acquisition, I'm imagining that recOrder will print out a waveorder CLI call to reproduce the result.

I realize this suggestion walks back one of our earlier decisions to make waveorder an arrays-only library, but it seems to me that having the waveorder library be arrays-only will be challenging to test since all of the end-to-end tests will require recOrder or a different visualization tool. If waveorder has a CLI then I can test end-to-end reconstructions in waveorder quite easily.

[mattersoflight]

Thanks @talonchandler for flagging this problem and coming up with a good solution. My notes:

1. This two-stage implementation that saves the expensive-to-compute OTFs as Zarr makes a lot of sense. The second step apply-inv-tf is embarrassingly parallel. Having it as a distinct CLI will allow us to use SLURM! (Slurm-scripts #337).
2. We discussed the structure of zarr store at some length offline. I lean towards storing all the calibration data (background images, OTFs) in a single zarr store.
3. I lean towards keeping these reconstruction CLI within recOrder so that all the HPC-related CLI is in one repository.
4. Alternative names for the stages: wo compute-otf and wo apply-inv-otf, or wo forward-model and wo inverse-algorithm.

If we cannot figure out a way to co-develop recOrder and waveOrder (e.g., invoke a recOrder CLI with breakpoint deep in waveOrder array-library), we should move all HPC-related CLI to waveOrder.

??? Each algorithm takes a different number of params,
# what's the best way to encode this information

Let's discuss this offline.

[talonchandler]

@mattersoflight @ziw-liu and I came to a few conclusions in an offline meeting:

• the CLI described here will live in recOrder
• we'll use some pydantic model inheritance for the "shared" parameters to simplify the user-facing names
• although the recOrder-waveorder interface is pretty direct, it's not direct enough for us to directly pass the parameter models to the waveorder calls because recOrder uses a subset of waveorder features. This means there will be some minimal parameter matching on the recOrder side, but nothing heavy. (e.g. 2D phase OTFs can take arbitrary defocus positions on the waveorder side but only equally-spaced defocus positions on the recOrder side, so recOrder will need to make a small translation).

[ieivanov]

This plan sounds good to me too. I would advocate for clear and descriptive, possibly longer, function and variable names so that there is less confusion for both users and future developers. For me, working with the dragonfly_automation code base has been a breeze largely because of Keith's (substantial) effort in code clarity and documentation. We should strive for getting at least half way there.

[talonchandler]

#347 and the issues it spawned complete this discussion, so I am closing this issue.