How to make arbitrary transpositions of associations and datasets

Question

Is there a general Transpose for Dataset/Association that will transpose arbitrary levels, such as

Transpose[array, perm]

which will permute the levels of array according to the permutation perm? One seems to be able to use Transpose[ds] to transpose levels 1 and 2 and Map[Query@Transpose, ds, {n}] to transpose levels n+1 and n+2. But they only work on "adjacent" levels. To transpose levels 1 and 3 (i.e. perm = {3, 2, 1}), I seem to have to do something like the following

Transpose[Transpose /@ Transpose@ assoc]
Transpose[ds][Transpose, Transpose]

Is there a general way to do this for an arbitrary permutation of the levels of the data?

(I give a solution below, but as I'm learning how to use Dataset, I figure there are better ways.)

Related: Association of Associations : how to permute level 1 and level 2 keys? and comments.

Answer 1

Definitions

Here is an alternative implementation using the Wolfram Function Repository functions
AssociationKeyFlatten and ToAssociations (submitted by WRI personnel) and the function meMerge (localMerge) from the answer by andre314:

Clear[TransposeAssoc];
TransposeAssoc[assoc_Association, perm_?PermutationListQ] :=
  Block[{assoc2, assoc3, LocalMerge},
   LocalMerge[x : {_Association ..}] := Merge[x, LocalMerge]; 
   LocalMerge[{x_}] := x;
   assoc2 = ResourceFunction["AssociationKeyFlatten"][assoc];
   assoc3 = KeyMap[Permute[#, perm] &, assoc2];
   LocalMerge[
    ResourceFunction["ToAssociations"]@
     KeyValueMap[Fold[{#2 -> #1} &, #2, Reverse@#1] &, assoc3]]
  ];

Step-by-step run

Answer 2

If you factor a permuation perm into a product of cycles of the form $(j\ k)$ with $k=j+1$, then the permuation can be effected by Query and Transpose.

Functions:

adjacentCycles[perm] (* factors perm into "adjacent" 2-cycles *)
dsTranspose[x, perm] (* like Transpose[x, perm],
                        but x is a Dataset or Association *)

Code:

(* adjacentCycles
 *   factor permutations into cycles of the form (n n+1)
 *)
adjacentCycles[p_?PermutationListQ] := 
  Flatten@iAdjacentCycles[PermutationCycles[p]];
adjacentCycles[c : Cycles[{{__Integer} ..}]] := 
  Flatten@iAdjacentCycles[c];
iAdjacentCycles[Cycles[c : {}]] := {};
iAdjacentCycles[Cycles[c : {c1_, c2__}]] :=(*Join@@*)
  iAdjacentCycles /@ Cycles@*List /@ c;
iAdjacentCycles[Cycles[{c : {x_, y_, z__}}]] :=(*Join@@*)
  iAdjacentCycles /@ Cycles@*List /@ Reverse@Partition[c, 2, 1];
iAdjacentCycles[Cycles[{c : {x_, y_}}]] := Module[{a, b},
   {a, b} = MinMax[{x, y}];
   With[{factors = 
      Cycles@*List /@ Reverse@Partition[Range[a, b], 2, 1]},
    Reverse@Rest[factors]~Join~factors]
   ];

ClearAll[dsTranspose];
dsTranspose[assoc_Association, perm_?PermutationListQ] := 
  With[{res = dsTranspose[Dataset@assoc, perm]},
   Normal@res /; Dataset`ValidDatasetQ[res]
   ];
dsTranspose[ds_Dataset, perm_?PermutationListQ] :=
  Module[{
    xps,  (* perm factored as 2-cycle transpositions *)
    xpFN, (* applies Transpose or Query[Transpose] to appropriate level *)
    res},
   xps = adjacentCycles@perm;
   xps = xps[[All, 1, 1, 1]] - 1; (* levels to be transposed *)
   xpFN[0] = Transpose;
   xpFN[n_Integer?Positive] := 
    Map[Check[Query[Transpose][#], 
        Throw[$Failed, dsTranspose]] &, #, {n}] &;
   res = Catch[Fold[xpFN[#2][#1] &, ds, xps], dsTranspose];
   res /; Dataset`ValidDatasetQ[res]
   ];

Example:

assoc = Fold[AssociationThread[#2 -> #1] &, "X", 
  Reverse@Table[ToString[10 i + j], {i, 4}, {j, 2}]]
(*
<|"11" ->
    <|"21" -> <|"31" -> <|"41" -> "X", "42" -> "X"|>, 
                "32" -> <|"41" -> "X", "42" -> "X"|>|>, 
      "22" -> <|"31" -> <|"41" -> "X", "42" -> "X"|>, 
                "32" -> <|"41" -> "X", "42" -> "X"|>|>|>, 
  "12" ->
    <|"21" -> <|"31" -> <|"41" -> "X", "42" -> "X"|>, 
                "32" -> <|"41" -> "X", "42" -> "X"|>|>, 
      "22" -> <|"31" -> <|"41" -> "X", "42" -> "X"|>, 
                "32" -> <|"41" -> "X", "42" -> "X"|>|>|>|>
*)

dsTranspose[assoc, {3, 1, 4, 2}]
(*                                     (* perm:         *)
<|"21" ->                              (*  level 2 -> 1 *)
    <|"41" ->                          (*  level 4 -> 2 *)
              <|"11" ->                (*  level 1 -> 3 *)
                        <|"31" -> "X", (*  level 3 -> 4 *)
                                       "32" -> "X"|>,
                "12" -> <|"31" -> "X", "32" -> "X"|>|>, 
      "42" -> <|"11" -> <|"31" -> "X", "32" -> "X"|>, 
                "12" -> <|"31" -> "X", "32" -> "X"|>|>|>, 
  "22" ->
    <|"41" -> <|"11" -> <|"31" -> "X", "32" -> "X"|>, 
                "12" -> <|"31" -> "X", "32" -> "X"|>|>, 
      "42" -> <|"11" -> <|"31" -> "X", "32" -> "X"|>, 
                "12" -> <|"31" -> "X", "32" -> "X"|>|>|>|>
*)

Answer 3

Here is an approach that doesn't need to flatten the whole tree of data.

Let's take the example of moving the level 1 keys at the level 3, on the data assoc of your self-answer :

assoc = Fold[AssociationThread[#2 -> #1] &, "X", 
  Reverse@Table[ToString[10 i + j], {i, 4}, {j, 2}]]

Here is a function showAssocListTree that will be usefull to show clearly how associations are nested. It's just a formatting function. It is useless to understand it:

showAssocListTree = RightComposition[
  # //. List[content___] :>  Prepend[List1 /@ List1[content], "List"] &
  , # /. List1 -> List &
  , # //. as : Association[___] :>  
     Prepend[List @@@ Normal[as], "Ass."] &
  , TableForm[#] &
  , ToBoxes
  , # //. GridBox[{{"\"List\"", ___}, r___}, r01___] :> 
     RowBox[{RotationBox[
        StyleBox["\"List\"", FontVariations -> {"Underline" -> True}],
         BoxRotation -> Pi/2], "["(*StyleBox["[",FontWeight\[Rule] 
       "Bold"]*), GridBox[{r}, r01]}] &
  , # //. GridBox[{{"\"Ass.\"", ___}, r___}, r01___] :> 
     RowBox[{"-> ", RotationBox["\"Ass.\"", BoxRotation -> Pi/2], 
       StyleBox["[", FontWeight -> "Bold"], GridBox[{r}, r01]}] &
  , # /. RowBox[{"-> ", r___}] :>  RowBox[{r}] &
  , # //. InterpretationBox[x_, ___] :> x &
  , # /. RowBox[{a___, RotationBox["\"Ass.\"", BoxRotation -> Pi/2], 
       r___}] :>  
     RowBox[{a, RotationBox["\"Association\"", BoxRotation -> Pi/2], 
       r}] &
  , RawBoxes
  , Style[#, GridBoxOptions -> { GridBoxDividers -> None}, 
    SpanMaxSize -> DirectedInfinity[1]] &
  ];

your data formatted :

assoc  // showAssocListTree

Insertion of level 1 data at level 3 :

listOfAssoc=KeyValueMap[
 Function[{k, v}, Map[Association[k -> #] &, v, {2}]], assoc];  

listOfAssoc //showAssocListTree 

Note that the outer Association has been transformed in a List.

Now, the built-in function Merge will be applied. This function only merge two successive levels. So, a recursive function is first created. Note that when applied, this recursive function will explode exponentially, but in most cases, it's better than to flatten the whole tree, and there are probably ways to circumvent this problem.

ClearAll[myMerge]
myMerge[x : {_Association ..}] := Merge[x, myMerge]
myMerge[{x_}] := x  


myMerge[listOfAssoc ] //showAssocListTree 

This approach can be generalised to the general case. For example {3, 1, 4, 2} could be decomposed as {1, 2, 3, 4} -> {1, 3, 4, 2} -> {3, 1, 4, 2}

inspiration source