[ENH] `TransformedDistribution` handling of index passing in subsetting

[fkiraly]

I noticed a few problems with TransformedDistribution subsetting, and it is an issue in-principle.

Comments on how to handle this would be appreciated.

Namely, consider a univariate td = TransformedDistribution(distr, trafo, inv_trafo).

Problem 1 - columns forgotten in scalar subsetting

Let's say trafo requires a specific column name to work.

If we do td_subset = td.iat[0,0], this produces a column-less scalar distribution.
However, whenever we pass data to trafo or inv_trafo, it will require the column name, which is no longer stored in td_subset in the current implementation.

It seems we need to remember the column name when we subset to scalar.

Problem 2 - multivariate transformations

Let's now consider td = TransformedDistribution(distr, trafo, inv_trafo) of shape (3, 2), i.e., two columns.

If we subset to a single column, applying trafo or inv_trafo to (3,1) objects no longer work, since both will require an (n, 2) object as per sklearn contract.

So, in the multivariate transformation case it is even worse, we seem to require a copied memory of the entire original distribution.

[fkiraly]

FYI @joshdunnlime - comments appreciated!

[joshdunnlime]

Problem 1
Having a DifferentiableTranformer in the backend would make it easy to solve these indexing issues. It retains the index and column attributes on the DT class, so in this case we can use the column attr to cast the scaler to a dataframe, pass it through the transformer, and then do to_scaler on the output.

Problem 2
For multivariate targets, we can also do something similar, cast to a dataframe with empty columns filled with NaNs, transform, then remove empty columns. This isn't pretty but it would work.

An alternative option would be using something like the sklearn FeatureUnion to do multivariate transformations in separate transformers. Then each transform in the union can be called individually. This probably requires more overhead though. It also means handling the "True Jacobian" (see point 5. here) will become trickier. That said, if we want to iat some multivariate distribution, that likely implies the target transforms can be treated as independent, that the Jacobian is diagonal and is then just the product of the non-partial derivatives.