Neural Updates: PolicyDataLoader replaced with the old dataloading scheme and Quad terms added

[selmanozleyen]

hi @MUCDK,

I have added support for quad terms while replacing the old dataloader with something more efficient and more extendable. Currently quad terms will work in a meaningful way if in the policy this is satisfied for the policy graph for all a,b,c,d in nodes, for every (a,b) in edges there is no (c,a) or no (b,d) present. I.e., a src node should stay a source node and tgt node should stay a tgt node throughout the plan. Currently I handled it like this

class NeuralOTProblem

    @wrap_prepare
    def prepare(
        self,
        policy_key: str,
        policy: Policy_t,
        lin: Mapping[str, Any],
        src_quad: Optional[Mapping[str, Any]] = None,
        tgt_quad: Optional[Mapping[str, Any]] = None,
        condition: Optional[Mapping[str, Any]] = None,
        subset: Optional[Sequence[Tuple[K, K]]] = None,
        seed: int = 0,
        reference: K = None,
    ) -> "NeuralOTProblem[K]":

Another problem is that I could've assume an adata has both src_quad and tgt_quad. For this we'd also need to assume there is no second adata. These assumptions can make it a bit easier but it depends on our main goal with adding quad support. For now it makes sense to me that all adata will have the aug and flow. This quad problem actually exists in non-Neural version right? How do we solve that?

How reasonable are these?

[MUCDK]

The assumption for all a,b,c,d in nodes, for every (a,b) in edges there is no (c,a) or no (b,d) present cannot hold, unfortunately. We often have the case of (a,b,c,d) with (a,b), (b,c), and (c,d), ii.e. sequential policy

[MUCDK]

Regarding the number of adatas, I think it's fine if we only allow for at most two adatas, but we must allow them to be different, espeically in the case of modality translation.

[selmanozleyen]

The assumption for all a,b,c,d in nodes, for every (a,b) in edges there is no (c,a) or no (b,d) present cannot hold, unfortunately. We often have the case of (a,b,c,d) with (a,b), (b,c), and (c,d), ii.e. sequential policy

But it should hold when we have only one adata and quadratic terms right? Because if we have two quadratic term attr's of one side that means we have it also on the other side. If we have the same quadratic attributes on both sides then it's reasonable to assume this should've been the linear term right since its in comparable space? In fact some rules like these might help us out for structuring:

Adata assumption: Every given adata will have attributes for all the cells. Ie it can't have empty entries for quadratic terms on some cells, if there is a term it should be present on all cells for that adata

1. For linear attr only and any number of adatas all policies make sense.

2. When there is at most 2 adatas given and quadratic attrs given. Only such policies make sense:
   - for all a,b,c,d in nodes, for every (a,b) in edges there is no (c,a) or no (b,d) present (i.e. no node seen as src will be seen as tgt and vice versa)
   - One might even argue that with quadratic term one adata might not even makes sense as the quadratic term is there for the whole adata

3. When more than 2 adatas and quadratic given (Optional):
   - There will be no subsetting of adatas i.e., every node will be one adata
   - Then all policies make sense

Recognizing these as rules would clarify the implementation a lot. From what I see so far we already don't violate these assumptions explicitly. Wdyt @MUCDK am I missing something?

[MUCDK]

Yes, all these points make sense!