Support asymmetry in the RIME

[sjperkins]

As Montblanc currently implements the RIME, symmetry in the left and right Jones terms are assumed. This is fast, computationally and makes optimal use of memory (especially on the GPU). The reasons for this are well-described in #249 and https://www.overleaf.com/read/ccvvkzgymntg#/55754428/.

In a general formulation, the RIME does not strictly require symmetry in the Jones terms. In particular, it would be useful to support different Direction Dependent and Independent Effect terms on the left and right side sides of the onion. In order to properly support this it will be necessary to handle the extra configuration and memory requirements.

There are two considerations:

1. User configuration of the left and right terms. This will be more completely dealt with in RIME configuration scheme #254.
2. Allowing a mechanism for the user to supply different inputs to the left and right terms.

With regard to 2., we probably don't want to support completely different sets as many inputs will be exactly the same for both sides of the onion (parallactic angles for example, which are expensive to compute) and maintaining two sets of input will be expensive in terms of memory.

On the other hand, we do want to support things like two cube inputs on either side of the onion. Beam cubes are occupy large amounts of memory and the memory budgeting mechanism will now need to account for two separate graph paths with corresponding inputs (beams) and outputs (ddes), for example. A proper and (probably fairly easy) solution would be to traverse the graph using tensorflow's shape inference mechanism with different chunk sizes to estimate the maximum memory usage.

With respect to providing inputs for the left and right sides, the intended interface for the dask version is something like

from xarrayms import xds_from_ms

ms_dataset = xds_from_ms("WSRT.MS")
mb_dataset = montblanc_dataset_from_ms(ms_dataset)

rime = montblanc.Rime(configuration={...})

model_vis = rime(mb_dataset)

The rime function should either take a single dataset (input for the symmetrical case) or a tuple composed of (left_dataset, everything_else_dataset, right_dataset)

So something like

model_vis = rime(right_beam_ds, primary_ds, left_beam_ds)

A slight contraction of the above could involve

model_vis = rime(primary_ds, left_beam_ds)

where right_beam_ds is folded into primary_ds

[sjperkins]

/cc @o-smirnov @landmanbester @JSKenyon