3
\$\begingroup\$

I'm working on some Python code and have a few functions which do similar things, and the only way I've found of writing them is quite ugly and not very clear.

In the example below, the goal is to compute the Kronecker product over a tensor chain of length M, in which the mth tensor is R and every other tensor is J.

Is there any nice way to rewrite this?

def make_rotate_target(m, M, J, R): 
    out = J
    if M == 1:
        return R
    else:
        for i in range(M):
            if i == 0:
                out = J
            else:
                if i + 1 == m:
                    out = np.kron(out, R)
                else:
                    out = np.kron(out, J)
    return out
\$\endgroup\$
1
2
\$\begingroup\$

functools.reduce is what you need here:

from functools import reduce

def make_rotate_target(m, M, J, R):
    input_chain = [J] * M
    input_chain[m - 1] = R
    return reduce(np.kron, input_chain)

The input_chain list could be replaced with an iterable constructed from itertools.repeat and itertools.chain to save space.

from functools import reduce
from itertools import repeat, chain

def make_rotate_target(m, M, J, R):
    input_chain = chain(repeat(J, m - 1), [R], repeat(J, M - m))
    return reduce(np.kron, input_chain)

The computation could be accelerated by exploiting the associative property of the Kronecker product: $$\underbrace{J\otimes J\otimes\cdots\otimes J}_{m-1}\otimes R\otimes\underbrace{J\otimes J\otimes\cdots\otimes J}_{M-m} \\ =(\underbrace{J\otimes J\otimes\cdots\otimes J}_{m-1})\otimes R\otimes(\underbrace{J\otimes J\otimes\cdots\otimes J}_{M-m}) $$ $$\underbrace{J\otimes J\otimes\cdots\otimes J}_{a + b}=(\underbrace{J\otimes J\otimes\cdots\otimes J}_{a})\otimes(\underbrace{J\otimes J\otimes\cdots\otimes J}_{b})$$

So some intermediate computation results could be reused. I'll leave the rest to you.

\$\endgroup\$
2
  • \$\begingroup\$ A beautiful solution. I knew something like this was possible. I think I'll be using the reduce function a lot. Thanks! \$\endgroup\$ – Dan Goldwater Sep 17 '20 at 14:57
  • 1
    \$\begingroup\$ @DanGoldwater FYI, many numpy functions (but not np.kron) has it own implementation of reduce (see doc) that supports reduction along an axis of an ndarray. You may need that in some cases. BTW, you could upvote my post while accepting it as an answer. \$\endgroup\$ – GZ0 Sep 17 '20 at 22:22
2
\$\begingroup\$

Assuming M is non-negative you can make it clear when J and R are returned and remove some of the nested if-else's.

Since if i+1 ==m only evaluates as true if m <= M you could check that and use different for loops if it's true/false to make things slightly faster but I think that would decrease readability for not much gain

def make_rotate_target(m, M, J, R):
    if M == 0:
        return J

    if M == 1:
        return R

    out = J
    for i in range(1, M):
        if i + 1 == m:
            out = np.kron(out, R)
        else:
            out = np.kron(out, J)

    return out
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.