A few months ago, during yet another of my hobby research side-projects (one I had a lot of fun with, but I eventually got stuck and gave up on it), I needed a general way to generate collections of distinct partitions of the same set with certain properties. I was not able to find any known and ready-to-implement algorithm, so I started to tinker, trying to come up with something. To my surprise, eventually I succeeded – but did not really understand why it works, beyond the glimpses of intuition that have lead me to my construction. I asked a question on the math StackExchange, which sadly went by completely unnoticed. Maybe I used the wrong terminology and keywords after all, or I should have asked elsewhere.

In any case, I decided to finally to write up what I discovered, before it gets lost forever in the sands of time and depths of my laptop. I will first introduce the problem, then I will explain my solution approach, and finally present some code.

Defining The Problem

First I will state the problem using basic terms, a bit later I will introduce and refer to relevant terminology and concepts that I found in literature.

A partition of some set $X$ is a set of disjoint sets (we will call blocks) whose union is $X$. Without loss of generality, we will only look at partitions of sets of the form $X_k := \{0, 1, …, k-1\}$ and sometimes call the elements of such a set points. A partition is uniform if all the blocks have the same size. For example, a uniform partition of $X_9$ is given by $\{\{0,5,7\}, \{1,6,8\}, \{2,3,4\}\}$. A partition is trivial if its blocks all just group a sequence of adjacent numbers together. The previous example is non-trivial, while $\{\{0,1,2\}, \{3,4,5\}, \{6,7,8\}\}$ is a trivial partition of $X_9$.

Two partitions are orthogonal if any two blocks intersect in at most one element. A set of partitions is orthogonal if all contained partitions are. For example, a set of uniform orthogonal set partitions of $X_9$ is given by $P = \{p_0, p_1, p_2, p_3\}$ with

• $p_0 = \{\{0,1,2\},\{3,4,5\},\{6,7,8\}\}$
• $p_1 = \{\{0,3,6\},\{1,4,7\},\{2,5,8\}\}$
• $p_2 = \{\{0,4,8\},\{1,5,6\},\{2,3,7\}\}$
• $p_3 = \{\{2,4,6\},\{1,3,8\},\{0,5,7\}\}$

Intuitively, these orthogonal partitions are “maximally unrelated”. Clearly, allowing the overlap of 1 element for blocks from distinct partitions is unavoidable – the partitions all use the same base set, and each element must appear in some block.

Now the problem: What I was looking for was a general method which takes a block size $b>1$ and $i>1$ as inputs, and for $n=b^i$ can generate as many uniform orthogonal partitions of $X_n$ with block size $b$ as possible.

Why would you want to do that? In my case, and without going into too much detail – I was trying to construct a certain kind of graph family, based on many copies of a gadget subgraph with fixed size $b$. I needed a way to ensure that nodes are members of multiple gadgets, but no two nodes are member of more than one common gadget. A set of orthogonal partitions provides a “skeleton” on which such a graph can be built by adding the desired edges. I hope that someone else maybe can also make some good use of this kind of object.

Doing My Homework: Some Hints From Literature

While researching how to generate any such orthogonal set partitions, I figured out that what I am looking for is connected to the well-studied field of combinatorial designs. I also discovered Mutually Orthogonal Latin Squares (MOLS), surprised by the amount of research and depth of the math connected to this kind of objects.

I found this code to easily generate a complete set of $\textrm{MOLS}(p)$, consisting of $p-1$ MOLS for a prime number $p$, and the Wikipedia article on MOLS also explains how to convert a set of MOLS into such partitions by translating them into Orthogonal Arrays, and finally into nets – which yield the desired partitions. This was a very good starting point.

Apparently there is also a field-based construction that allows to get all $p^k-1$ MOLS for prime powers $p^k$, but in order to effectively construct a representation of that field and work with it, one needs certain suitable polynomials which are in general difficult to find. So from a practical point of view, this is not very useful.

I also stumbled on the social golfer problem, which finally led me to the notion of resolved group divisible designs (RGDD) – that concept looks like exactly the kind of definition I was looking for. Now, to state my problem using the terminology from the linked paper: what I was looking for are RGDDs with $\lambda=1$ (i.e., blocks from different partitions intersect at most in 1 element), and more precisely, families of RGDDs with fixed and singleton set $K$ (i.e., uniform block size), but variable $\nu$ (i.e., number of elements in the underlying set).

In general, block designs are not strongly related to set partitions, it is a much more general concept. But a resolved block design is defined as a separation of the blocks into groups (also called resolutions) such that each group is a partition of the underlying set. A resolved block design is group divisible if the groups are orthogonal, i.e. blocks of the same group do not intersect and blocks between groups intersect in at most one element. Hence, the notion of RGDDs is able to express exactly what is needed here.

As, sadly, my aforementioned StackExchange question (which this post has basically just been following until now) received no replies, I got stuck at this point of researching and understanding the theory, but at least I found some relevant terms and concepts.

Theory: How to Pluck Set Partitions From Trees

The construction I came up with takes a “base RGDD” (e.g. obtained from MOLS, or by any other means), and basically “lifts” it to larger sets, i.e. larger values of $n$. To be precise – for a fixed block size $b > 1$, given a set of $k$ distinct non-trivial uniform partitions on $b^2$ elements with block size $b$, the construction will produce for any $i > 1$ a collection of $k^{i−1}$ distinct uniform mutually orthogonal set partitions of $b^i$ elements, each consisting of $b^{i−1}$ blocks with size $b$.

Creative Leaps: From Sets to Lists to Matrices to Trees

We will use the example set of partitions $P$ on $X_9$ you have seen above as a running example. We can use it to form a base RGDD with $b=3, k=3, n=b^2=9$. As a reminder: $b$ is the block size, $k$ is the number of partitions and $n$ the number of points being partitioned. Out of the four partitions in $P$, the first one, $p_0 = \{\{0,1,2\},\{3,4,5\},\{6,7,8\}\}$, is the unique trivial one we must exclude. So in this case, our base RGDD for $b=3$ is the set $\Sigma := \{p_1, p_2, p_3 \} = P \setminus \{p_0\}$.

Let us pretend that a partition is not a set of sets, but has order, i.e. is a list of lists¹. As for the base partition set we have $n=b^2$, we can also write them as square matrices:

$$ \def\arraystretch{1.5} p_0 = \begin{array}{|c:c:c|} \hline 0 & 1 & 2 \\ \hline 3 & 4 & 5 \\ \hline 6 & 7 & 8 \\ \hline \end{array} \hspace{3mm} p_1= \begin{array}{|c:c:c|} \hline 0 & 3 & 6 \\ \hline 1 & 4 & 7 \\ \hline 2 & 5 & 8 \\ \hline \end{array} \hspace{3mm} p_2= \begin{array}{|c:c:c|} \hline 0 & 4 & 8 \\ \hline 1 & 5 & 6 \\ \hline 2 & 3 & 7 \\ \hline \end{array} \hspace{3mm} p_3= \begin{array}{|c:c:c|} \hline 0 & 5 & 7 \\ \hline 1 & 3 & 8 \\ \hline 2 & 4 & 6 \\ \hline \end{array} $$

A matrix can be obviously interpreted (or represented) as a sequence of rows, which matches the typical mathematical indexing. We can also interpret the matrix as a flat list of numbers, if we just concatenate the rows – this is another typical representation of matrices, often used in implementations. So given a $n\times n$ matrix $M$, the element $M[r,c]$ in row $r$ and column $c$ corresponds to the index $n*r + c$ in the flattened list. So far, so well-known.

Now for something more interesting – take a set partition, in its matrix representation, and flatten it. What we got is a list of all points in $X_{n^2}$, each appearing exactly once. What could this possibly be? It is a permutation! It means that these matrices can be interpreted as permutations of cells in other matrices of the same size. For example, the trivial partition $p_0$ corresponds to the trivial identity permutation, i.e. leaves a matrix unchanged², while $p_1$ corresponds to a matrix transposition.

How can we possibly use these permutations to act on more than just matrices of the same shape? Time to get creative. What if instead of values the cells would contain another sequence of values? And in fact, why stop there and not allow to repeat this “nesting” up to any depth? You probably can see where this is going – we need to move away from matrices and use trees instead.

An $n\times n$ matrix can be interpreted as a $n$-regular ordered tree of depth $2$:

• the root node is the whole matrix
• the first level of children are the rows
• the second level are the cells in the respective row

For example, the matrix for $p_1$ looks like this as a tree:

The key idea now is to consider regular ordered trees of arbitrary depth. In such a tree, the $n$ leaves are uniquely labeled by a point from the underlying set $X_n$ (recall that in general we assume $n=b^i$ for some $i$), and each group of leaves sharing the parent node defines a block of the partition the tree is representing. We will call these leaf-labelled trees the index trees on $X_n$. To each partition of $X_{n}$ with uniform block size $b$ we can associate a canonical index tree, similarly to how we interpreted partitions as matrices. The standard index tree then is the canonical index tree corresponding to the unique trivial partition on $X_{n}$.

If you wonder why I call them index trees – there is another intuitive and geometric way to think about these trees: an $n$-regular tree of depth $i$ can be seen as an indexing scheme into an $i$-dimensional hypercube with side length $n$. If the leaves are points, then their parents are lines, whose parents are planes, whose parents are volumes, and so on. In a sense, this view is closer to the actual implementation³, but we will mostly stick with the tree view.

If we want to partition a larger set, we can simply use a tree with more levels (i.e. use a higher-dimensional hypercube). For example, this is the standard index tree for $n = b^3 = 27$ elements:

Luckily, permutations do not care what they permutate. We have seen a simple correspondence between trees of depth 2 and matrices. If the tree that our permutation acts on happens to have a greater depth, we will just permutate the subtrees. If we apply $p_1$ to this tree, we will get:

Recall that $p_1$ was just matrix transposition, so what we did was essentially this:

$$ \footnotesize \def\arraystretch{1.5} \begin{array}{|c:c:c|} \hline (0,1,2) & (3,4,5) & (6,7,8) \\ \hline (9,10,11) & (12,13,14) & (15,16,17) \\ \hline (18,19,20) & (21,22,23) & (24,25,26) \\ \hline \end{array} \hspace{2mm} \overset{p_1}{\rightarrow} \hspace{2mm} \begin{array}{|c:c:c|} \hline (0,1,2) & (9,10,11) & (18,19,20) \\ \hline (3,4,5) & (12,13,14) & (21,22,23) \\ \hline (6,7,8) & (15,16,17) & (24,25,26) \\ \hline \end{array} $$

But as our tree is regular, the same trick that works at the root also works on any subtree of depth (at least) 2! For example, what we can do is take a different permutation, say $p_2$, and apply it to the first subtree:

Now you have seen all the main ingredients that we need, so we are ready to put them together and get the promised construction.

Twisted Trees: It’s Permutations All The Way Down

Recall that we have the base RGDD on $b^2$ points with $k$ non-trivial partitions, each with block size $b$, and know how to interpret them as permutations. We also have a method to apply such a permutation to any $b$-regular (sub-)tree that has depth $\geq 2$, and this operation will result in a modified tree where the grandchild-nodes (i.e., the subtrees rooted at these nodes) are reordered.

What we can do once, we can also do multiple times. So given a $b$-regular tree of depth $i$, we will do the following: pick a sequence of $(i-1)$ permutations from the base RGDD (they are allowed to repeat). Then, go down level by level, and apply the corresponding permutation to each node (i.e. subtree) at that level. For example, recall that we applied $p_1$ at the root (level 0) of the standard index tree on $X_{27}$, and then we applied $p_2$ to one of the subtrees (level 1). Now what we have to do is apply $p_2$ to the remaining level 1 subtrees as well, yielding the following tree:

We started with the standard index tree which corresponds to the trivial partition of $X_{27}$. We traversed all levels we can permutate (i.e., where nodes have a depth of at least 2) and applied the permutation chosen for that level (we have selected the sequence of permutations $p_1, p_2$). Now we can harvest the fruit of our labour by reading the set partition blocks directly off the leaves from resulting tree:

$$ \begin{Bmatrix} \{0, 10, 20\} & \{1, 11, 18\} & \{2, 9, 19\} \\ \{3, 13, 23\} & \{4, 14, 21\} & \{5, 12, 22\} \\ \{6, 16, 26\} & \{7, 17, 24\} & \{8, 15, 25\} \end{Bmatrix} $$

We have used one possible sequence (of length 2) of non-trivial permutations on our tree (of depth 3), and thus got one possible new partition. How do we get others? Well, we have 3 suitable permutations $\Sigma = \{p_1, p_2, p_3\}$, so there are $3^2$ sequences $w \in \Sigma^2$ of the required length we can form - if we perform the same procedure for each such permutation sequence $w$, every time starting from the standard index tree, then for each $w$ we will get a distinct partition, and in fact it turns out, all the resulting partitions are mutually orthogonal.⁴

Seeding the Tree: Computing a Base RGDD

After you have seen how the construction works, let us briefly discuss how to get a suitable set of partitions to get started. A simple way to get a suitable base RGDD is using the MOLS-based method I mentioned earlier. If $b$ is prime, compute the complete set of all $b-1$ distinct $\mathrm{MOLS}(b)$ (i.e., for squares with side length $b$), convert these to $b+1$ orthogonal arrays, extract the nets (which are our partitions), and discard the unique trivial partition, leaving you with $b$ partitions that are usable in the construction (so with this method we can ensure that $k=b$, which is neat).⁵

But I found a simpler, direct way to get the same result⁶. We do not actually have to go through all these steps, compute MOLS, transform them back and forth, only to get our hands on some nice set partitions.

Recall our set of partitions $P$, which is based on the complete set of $\mathrm{MOLS}(3)$. Let us look again very closely at those partitions, written in matrix form, and paying special attention to the columns:

$$ \def\arraystretch{1.5} p_0 = \begin{array}{|c:c:c|} \hline 0 & 1 & 2 \\ \hline 3 & 4 & 5 \\ \hline 6 & 7 & 8 \\ \hline \end{array} \hspace{3mm} p_1= \begin{array}{|c|c|c|} \hline 0 & 3 & 6 \\ 1 & 4 & 7 \\ 2 & 5 & 8 \\ \hline \end{array} \hspace{3mm} p_2= \begin{array}{|c|c|c|} \hline 0 & 4 & 8 \\ 1 & 5 & 6 \\ 2 & 3 & 7 \\ \hline \end{array} \hspace{3mm} p_3= \begin{array}{|c|c|c|} \hline 0 & 5 & 7 \\ 1 & 3 & 8 \\ 2 & 4 & 6 \\ \hline \end{array} $$

Can you already see it? Let me also show you the matrices for $b=5$, obtained from the MOLS-based set partition construction:

$$ \def\arraystretch{1} q_0 = \begin{array}{|c:c:c:c:c|} \hline \bf 0 & 1 & 2 & 3 & 4 \\ \hline \bf 5 & 6 & 7 & 8 & 9 \\ \hline \bf 10 & 11 & 12 & 13 & 14 \\ \hline \bf 15 & 16 & 17 & 18 & 19 \\ \hline \bf 20 & 21 & 22 & 23 & 24 \\ \hline \end{array} \hspace{3mm} q_1 = \begin{array}{|c|c|c|c|c|} \hline \bf 0 & 9 & 13 & 17 & 21\\ 1 & \bf 5 & 14 & 18 & 22\\ 2 & 6 & \bf 10 & 19 & 23\\ 3 & 7 & 11 & \bf 15 & 24\\ 4 & 8 & 12 & 16 & \bf 20\\ \hline \end{array} \hspace{3mm} q_2 = \begin{array}{|c|c|c|c|c|} \hline \bf 0 & 8 & 11 & 19 & 22\\ 1 & 9 & 12 & \bf 15 & 23\\ 2 & \bf 5 & 13 & 16 & 24\\ 3 & 6 & 14 & 17 & \bf 20\\ 4 & 7 & \bf 10 & 18 & 21\\ \hline \end{array} $$

$$ \def\arraystretch{1} \hspace{2mm} q_3 = \begin{array}{|c|c|c|c|c|} \hline \bf 0 & 7 & 14 & 16 & 23\\ 1 & 8 & \bf 10 & 17 & 24\\ 2 & 9 & 11 & 18 & \bf 20\\ 3 & \bf 5 & 12 & 19 & 21\\ 4 & 6 & 13 & \bf 15 & 22\\ \hline \end{array} \hspace{3mm} q_4 = \begin{array}{|c|c|c|c|c|} \hline \bf 0 & 6 & 12 & 18 & 24\\ 1 & 7 & 13 & 19 & \bf 20\\ 2 & 8 & 14 & \bf 15 & 21\\ 3 & 9 & \bf 10 & 16 & 22\\ 4 & \bf 5 & 11 & 17 & 23\\ \hline \end{array} \hspace{3mm} q_5 = \begin{array}{|c|c|c|c|c|} \hline \bf 0 & \bf 5 & \bf 10 & \bf 15 & \bf 20\\ 1 & 6 & 11 & 16 & 21\\ 2 & 7 & 12 & 17 & 22\\ 3 & 8 & 13 & 18 & 23\\ 4 & 9 & 14 & 19 & 24\\ \hline \end{array} $$

The $b$ non-trivial matrices corresponding to the non-trivial partitions are all very regular, and in fact can be easily obtained by:

• transposing the matrix of the trivial partition, i.e. taking a copy of ${q_0}^{\top}$
• cyclically rotating the values in column $j$ of $q_i$ by $i*(j-1)$ positions

Keep in mind that the rows are, as before, the blocks of the resulting partitions. Also note that using this approach, the partition corresponding to matrix transposition (i.e. just like $p_1$) will always be the last one in the sequence (i.e. in this case $q_5$). At this point, the column rotations become trivial, because the shift is a multiple of the column size. I think this approach is a neat shortcut, but note that just like the simple MOLS-based approach, it works only correctly for values of $b$ that are prime.

But in the end, it does not matter how the base RGDD was obtained, as long as you can somehow provide $k\geq 2$ distinct, non-trivial partitions on $b^2$ with block size $b$ we can work with.

Wrapping It Up: The Complete Construction

After we have seen all the steps in a (hopefully instructive) example for the parameters $b=k=i=3$, and hoping that I could provide some intuition for the construction, here is a brief summary of the full procedure:

• obtain a base RGDD $\Sigma$ that consists of $k\geq 2$ non-trivial partitions of $b^2$ points, each with uniform block size $b\geq 2$ (converted from MOLS, or otherwise)
• pick the target set size $b^i$ with $i\geq 2$ and construct the standard index tree $t_0$
• for each permutation sequence $w=w_0w_1… \in \Sigma^{i-1}$, going level by level, apply the permutation $w_i$ to each node on level $i$ of $t_0$, yielding a new tree $t_w$
• flatten the resulting trees $t_w$ into partitions, i.e. sibling leaves form a block

You now should have $k^{i-1}$ distinct, non-trivial, mutually orthogonal set partitions of $b^i$ points, each with $b^{i-1}$ blocks of size $b$. You can find an implementation of all this further below.

But Why Does It Work? Some Ideas and Guesses

Concerning the MOLS-free approach to generate the initial partitions, I have the suspicion that it probably works for similar reasons to why the MOLS can be easily generated using index shifts in one line of code, probably implicitly exploiting the primality of $b$. To prove it, maybe one can use some similar number theoretic arguments. One could also attempt to show a correspondence of the constructions, but this is probably pretty ugly. This is all I can comment on that part, and in fact I discovered this shortcut while writing this post, so I have not thought much about it.

The central question is, of course, concerning the correctness of the lifting construction. From the explanations above it should be clear that the construction produces uniform partitions of the desired shape (because all we ever do is move things around). But how can we be sure that the resulting partitions are really mutually orthogonal, as claimed?

While I am quite confident that it works, the way I came up with it all was just playful exploration of the representational ambiguity of different objects⁷. You have seen all the ideas leading to it, and to me this construction just somehow feels right, but without a proof this of course does not have to mean anything at all.

I don’t see an immediately obvious and clear path towards proving it, and because this was just a sub-problem of something else I have been working on, it did not really matter to me as long as it did the job – which I just confirmed empirically. At least, the construction is simple enough for me to not expect any rare and hard-to-trigger edge cases – if it manages to correctly generate $5^3=125$ mutually orthogonal partitions of $5^4=625$ elements, I strongly believe it will work for any valid input parameters.

Even though I did not try to do it, I feel like the ideas and intuitions behind the construction can somehow be used, with some more careful analysis, to extract a correctness proof using some ad-hoc structural induction on the tree depth, if one identifies a few invariants about all the “tree shuffling” that is going on.

Yet, as the construction heavily relies on permutations, maybe what is actually happening can be explained elegantly using some group theory. After some superficial research I found that the concept of wreath products looks (at least superficially) relevant, because I compose the same set of permutations with itself, over and over. Maybe what I am doing here is equivalent to some obscure and simple lemma in a standard group theory textbook, I don’t know. But I believe thinking in terms of group actions might be fruitful here⁸.

Also, by Theorem 2.2 of the mentioned RGDD paper, it looks like what I am doing could possibly be an implementation for an iterated RGDD multiplication. At least it matches with the expected parameters of the RGDD according to the theorem, when plugging in the parameters $k = g = b, \lambda = \mu = 1, v = b^i, m = b^{i-1}$. My repeated use of the base RGDD seems to correspond to the role of the RGDD denoted by $GD[k, \mu, g, k_1g]$ in the theorem. As I apply the multiplication repeatedly, in fact I actually do exponentiation – that is another reason why it feels natural to refer to the initial partition set as the base RGDD.

But all of this is just my speculation about what might be going on. If anyone finds this interesting, maybe has some deeper insight to share, or even has time and motivation to prove this or investigate it further⁹, feel free to get in touch! If this is something actually new, I would love to eventually put it at least on arXiv, but as it stands, it is not ready to be put anywhere else but a blog post.

Practice: Stop the Talking and Show Some Code!

If you have read this whole post up to this point – thanks for your time and interest! As a reward, here I provide the promised implementation of the construction in Python. You can use it if you want to generate such orthogonal partitions for some reason, or just to play with it and test everything for yourself. Even though this was a pretty long post to explain what is going on (longer than I anticipated it would be), the implementation is actually very straight-forward once you know what to do.

For the base RGDD, the following code implements both the known MOLS-based approach and the (unproven) direct method I found. I also added some (inefficient) helper functions to verify the output. But without the MOLS and testing functions, it is actually just 50 lines of code, including comments.

"""Code to generate orthogonal set partitions with a uniform block size."""
# Copyright (C) 2023 Anton Pirogov, licensed under the MIT License
from itertools import product, repeat
import numpy as np

def prime_mols(n: int):
    """Generate complete set of MOLS for prime number n."""
    return [[[(k*i + j) % n for j in range(n)] for i in range(n)] for k in range(1, n)]

def mols_to_mops(mols):
    """Convert MOLS into orthogonal set partitions via orthogonal arrays.

    Returns n+2 set partitions of n^2 elements from a set of n MOLS.
    The returned list will always start with the trivial partition
    (i.e., counting up in sequential order, separated into groups).

    See https://en.wikipedia.org/wiki/Mutually_orthogonal_Latin_squares#Nets
    """
    oa = [  # convert to OA
        [r, c] + [mols[i][r][c] for i in range(len(mols))]
        for r in range(len(mols[0]))
        for c in range(len(mols[0][0]))
    ]
    num_mols, ls_sz, oa_rows = len(mols), len(mols[0]), len(oa)
    rc_j = [ # partitions based on rows/cols of OA
        [
            [ ls_sz * oa[k][0] + oa[k][1]
              for k in range(oa_rows) if oa[k][row_or_col] == j
            ] for j in range(ls_sz)
        ] for row_or_col in range(2)
    ]
    l_ij = [ # partitions induced by latin square columns
        [
            [
                ls_sz * oa[k][0] + oa[k][1]
                for k in range(oa_rows) if oa[k][2 + i] == j
            ] for j in range(ls_sz)
        ] for i in range(num_mols)
    ]
    return rc_j + l_ij

def base_rgdd(n: int):
    """Return a base RGDD for prime number n.

    The RGDD consists of n mutually orthogonal, non-trivial partitions
    of n^2 elements, each with uniform block size n.
    """
    return np.array(mols_to_mops(prime_mols(n))[1:])

# ----

def base_rgdd2(n: int):
    """Like base_rgdd, but direct and avoiding MOLS construction."""
    std = np.arange(n**2).reshape(n, n)
    return np.array([
        np.transpose(np.array([np.roll(std[i], k*i) for i in range(n)]))
        for k in range(1, n+1)
    ])

def get_tree(b: int, i: int):
    """Get standard index tree for b^i elements."""
    return np.arange(b**i).reshape(*repeat(b,i))

def to_partition(tree):
    """Convert an index tree into a set partition."""
    return tree.reshape(-1, tree.shape[-1])

def permutate_subtrees(tree, lvl: int, perm_mat):
    """Apply a permutation (given as a matrix) to all tree nodes at some level."""
    dim, deg = tree.ndim, len(tree)
    old_shape = list(repeat(deg, dim))# save old shape
    # build numpy indexing expression for subtree permutation
    perm_expr = list(repeat(slice(None),lvl)) + [perm_mat]
    # target shape = "flatten" the currently focused tree level
    flat_lvl = list(repeat(deg, lvl)) + [deg**2] + list(repeat(deg, dim - lvl - 2))
    # apply permutation in terms of numpy operations
    return tree.reshape(*flat_lvl)[*perm_expr].reshape(old_shape)

def permutate_tree(tree, perm_word):
    """Apply a sequence of permutations to an index tree."""
    assert tree.ndim == len(perm_word)+1
    for lv, lv_perm in enumerate(perm_word):
        tree = permutate_subtrees(tree, lv, lv_perm)
    return tree

def exp_partitions(perm_mats, i: int):
    """Lift a set of non-trivial orthogonal partitions on b^2 points to b^i points."""
    assert i > 1
    assert perm_mats[0].ndim == 2 and perm_mats[0].shape[0] == perm_mats[0].shape[1]
    block_size = perm_mats[0].shape[1]
    std_tree = get_tree(block_size, i)

    perms = [p.reshape(-1).tolist() for p in perm_mats]
    perm_words = product(*repeat(perms, i-1))
    return [to_partition(permutate_tree(std_tree, word)) for word in perm_words]

def gen_partitions(b: int, i: int):
    """For a prime b, generate b^(i-1) orthogonal partitions of b^i points."""
    return exp_partitions(base_rgdd2(b), i)

# ----

def partitions_orthogonal(p1, p2) -> bool:
    """Check whether two partitions are orthogonal."""
    return all(map(lambda b: len(set(b[0]).intersection(set(b[1]))) <= 1, product(p1, p2)))

def is_uniform_partition(p, b_size=None) -> bool:
    """Check that p is a partition of numbers 0 upto a value with uniform block size."""
    b_size = b_size or len(p[0])
    ns = p.reshape(-1).tolist()
    is_partition = set(ns) == set(range(len(ns)))  # is partition on 0..max
    is_uniform = all(map(lambda b: len(b) == b_size, p))  # uniform block size
    return is_partition and is_uniform

def sym_pairs(*vs):
    """Return all unordered pairs (i.e. outputs 1,2 but not 2,1)."""
    s = list(vs)
    return ((s[i], s[j]) for i in range(len(s)) for j in range(i+1, len(s)))

def check_partitions(ps):
    """Check that the computed partitions are uniform and mutually orthogonal."""
    b_size = len(ps[0][0])
    assert all(map(lambda p: is_uniform_partition(p, b_size), ps))
    for i, j in sym_pairs(*range(len(ps))):
        assert partitions_orthogonal(ps[i], ps[j])

Usage example:

>>> base_rgdd2(3)
array([[[0, 5, 7],
        [1, 3, 8],
        [2, 4, 6]],

       [[0, 4, 8],
        [1, 5, 6],
        [2, 3, 7]],

       [[0, 3, 6],
        [1, 4, 7],
        [2, 5, 8]]])
>>> gen_partitions(3, 3)
[
    array([[ 0, 17, 22], [ 1, 15, 23], [ 2, 16, 21], [ 3, 11, 25], [ 4,  9, 26], [ 5, 10, 24], [ 6, 14, 19], [ 7, 12, 20], [ 8, 13, 18]]),
    array([[ 0, 16, 23], [ 1, 17, 21], [ 2, 15, 22], [ 3, 10, 26], [ 4, 11, 24], [ 5,  9, 25], [ 6, 13, 20], [ 7, 14, 18], [ 8, 12, 19]]),
    array([[ 0, 15, 21], [ 1, 16, 22], [ 2, 17, 23], [ 3,  9, 24], [ 4, 10, 25], [ 5, 11, 26], [ 6, 12, 18], [ 7, 13, 19], [ 8, 14, 20]]),
    array([[ 0, 14, 25], [ 1, 12, 26], [ 2, 13, 24], [ 3, 17, 19], [ 4, 15, 20], [ 5, 16, 18], [ 6, 11, 22], [ 7,  9, 23], [ 8, 10, 21]]),
    array([[ 0, 13, 26], [ 1, 14, 24], [ 2, 12, 25], [ 3, 16, 20], [ 4, 17, 18], [ 5, 15, 19], [ 6, 10, 23], [ 7, 11, 21], [ 8,  9, 22]]),
    array([[ 0, 12, 24], [ 1, 13, 25], [ 2, 14, 26], [ 3, 15, 18], [ 4, 16, 19], [ 5, 17, 20], [ 6,  9, 21], [ 7, 10, 22], [ 8, 11, 23]]),
    array([[ 0, 11, 19], [ 1,  9, 20], [ 2, 10, 18], [ 3, 14, 22], [ 4, 12, 23], [ 5, 13, 21], [ 6, 17, 25], [ 7, 15, 26], [ 8, 16, 24]]),
    array([[ 0, 10, 20], [ 1, 11, 18], [ 2,  9, 19], [ 3, 13, 23], [ 4, 14, 21], [ 5, 12, 22], [ 6, 16, 26], [ 7, 17, 24], [ 8, 15, 25]]),
    array([[ 0,  9, 18], [ 1, 10, 19], [ 2, 11, 20], [ 3, 12, 21], [ 4, 13, 22], [ 5, 14, 23], [ 6, 15, 24], [ 7, 16, 25], [ 8, 17, 26]])
]
>>> check_partitions(gen_partitions(3, 3))  # no output = no errors
>>>