chain, chain.from_iterableの紹介(pythonのitertoolsを使いこなすために) #Python

前回の記事 isliceの紹介、具体例添えに続いて、itertools を使いこなすために、関数とその具体的な利用方法を紹介する。

今回はchain (とchain.from_iterable) の紹介をする。

chainとは

先頭の iterable の全要素を返し、次に2番目の iterable の全要素を返し、と全 iterable の要素を返すイテレータを作成します。連続したシーケンスを一つのシーケンスとして扱う場合に使用します。
itertools.chain 公式ドキュメントより

簡単に言えば、chainは(鎖という文字の意味通り)複数のiterableをつなぐことができる。

Pythonは動的型付けの言語 ¹ ではあるが、基本的には違う型同士をキャストなしにつなぐことはできない。

例えば、listとtupleは + でつなぐことができない。

sample_chain1.py

a = [1, 2, 3]
b = (4, 5, 6)
a + b
>> TypeError: can only concatenate list (not "tuple") to list

しかし、chainをつかうことで、iterableであれば、違う型を繋いで、ひとつのiteratorにできる。
もちろん、キャストを使えばlistやtupleにもできる。

sample_chain2.py

from itertools import chain

chain(a,b)
>>  <itertools.chain at xxxxxxxxxxx>
list(chain(a,b))
>> [1, 2, 3, 4, 5, 6]
tuple(chain(a,b))
>> (1, 2, 3, 4, 5, 6)

chain と chain.from_iterable

chainの亜種として、chain.from_iterableがある。

chain() のためのもう一つのコンストラクタである。
from chain.from_iterable 公式ドキュメントより

二つの大きな違いは受け取る引数でわかる。

chainは itertools.chain(*iterables) なのでアンパックした引数をうけとる
chain.from_iterable は chain.from_iterable(iterable) なので、引数は ひとつ しか受け取れない
- ソースコード(cpython) を読むと、chain.from_iterableは、受け取ったiterableを *iterables に書き換えて chain に渡している

参考 What is the difference between chain and chain.from_iterable in itertools?

具体例

実際にchainがどう使われているか調べてみた。

型が異なるiterableを繋げることができるので、iteratorやgeneratorが返ってくる関数の返り値をひとつにまとめるために使われることが多い印象がある。
また、連続したシーケンスを一つのシーケンスとして扱う という側面から、箱入れ的に深さのある要素 (list内のlistやndarrayなど) を平らな要素に変換する役割として使われることもあるようだ。

複数の引数について一度で処理 (jinja)

compiler.py

        kwarg_workaround = False
        for kwarg in chain((x.key for x in node.kwargs), extra_kwargs or ()):
            if is_python_keyword(kwarg):
                kwarg_workaround = True
                break

from jinja/jinja2/compiler.py

keyword argument (kwargs) が pythonですでに使われているキーワードでないかkeyword.iskeyword (コード中の is_python_keyword )を利用して確認するコード。
nodeのkwargsだけでなく、extra_kwargs も一度に確認するためにchainでつないでいる。つまり、以下と等価。

        kwarg_workaround = False
        for kwarg in (x.key for x in node.kwargs):
            if is_python_keyword(kwarg):
                kwarg_workaround = True
                break

        if(extra_kwargs):
                for kwarg in extra_kwargs:
                        if is_python_keyword(kwarg):
                                kwarg_workaround = True
                                break

行列の要素を繋げる (chainer)

n_step_rnn.py

def cudnn_rnn_weight_concat(
        n_layers, states, use_bi_direction, rnn_mode, ws, bs):
    rnn_dir = 'bi' if use_bi_direction else 'uni'
    inputs = itertools.chain(
        itertools.chain.from_iterable(ws),
        itertools.chain.from_iterable(bs),
    )
    return CudnnRNNWeightConcat(n_layers, states, rnn_dir, rnn_mode)(*inputs)

from chainer/functions/rnn/n_step_rnn.py

cuDNNのAPIに従ったweight matrices を作成するための、中間処理を行う関数。
chainが ws (Weight matrices) と bs (Bias vectors) とをつないでいる。wsとbsは chainer.Variable のオブジェクトであり、多次元配列である。
よって inputs は ws と bs の要素全ての要素を返すiteratorである。

多次元配列のchainとchain.from_iterableの挙動

このコードでは、多次元配列に対してchainとchain.from_iterableを併用している。

itertools.chain(itertools.chain.from_iterable(ws), itertools.chain.from_iterable(bs)) では何が起こっているのか？

実際に numpyのndarrayで実行してみる。

from itertools import  chain
import numpy as np

nums = np.array([[0, 1, 2], [10, 11, 12], [20, 21, 22]])
zeros = np.zeros((2, 3))

nums と zeros はそれぞれ(2, 3) の行列として宣言した。

list(chain(nums, zeros))
>>[array([0, 1, 2]), array([10, 11, 12]), array([20, 21, 22]), array([0., 0., 0.]), array([0., 0., 0.])]

まず、 chain(nums, zeros) の場合、以下の順番でiteratorを返す。

numsの中身 (nums[0], nums[1])
zerosの中身 (zeros[0], zeros[1])

このため、arrayのiteratorを返すことになる。

list(chain.from_iterable(nums))
>> [0, 1, 2, 10, 11, 12, 20, 21, 22]

一方、chain.from_iterable(nums) は配列の中身ひとつひとつを返すので、arrayではなく、要素をひとつづつ順番に返す。
つまり

nums[0][0], nums[0][1], ..., nums[1][2]

を順番に返す。よって多次元配列の要素を平らにしたiteratorとなる。

list(chain(chain.from_iterable(nums), chain.from_iterable(zeros))
>> 0, 1, 2, 10, 11, 12, 20, 21, 22, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]

最後に、chain(chain.from_iterable(nums), chain.from_iterable(zeros) は、 numsの要素、zeros の順番に、要素をひとつひとつ繋げて返すiteratorになる。

つまり、 itertools.chain(itertools.chain.from_iterable(ws), itertools.chain.from_iterable(bs)) は ws と bs の要素をひとつひとつ繋げて返すiteratorである。

map返り値の列挙 (magnitude)

converter_shared.py

def char_ngrams(key, beg, end):
    return chain.from_iterable((imap(lambda ngram: ''.join(ngram), zip(
        *[key[i:] for i in xrange(j)])) for j in xrange(beg, min(len(key) + 1, end + 1))))

from magnitude/pymagnitude/converter_shared.py

keyの文字ngram (beggram ~ endgram) を返す関数。
importまわりが特殊で、ここだけ引用するとわかりにくい。読みやすいよう、少し修正すると以下のようになる。

def char_ngrams(key, beg, end):
    return chain.from_iterable(
        (
            map( lambda ngram: ''.join(ngram), zip(*[key[i:] for i in range(j)]) )
            for j
            in range(beg, min(len(key) + 1, end + 1))
        )
    )

簡単に説明すると、関数内では以下の動作が行われている。

in range~ の部分で、beg ~ end の範囲を列挙
map(lambda ~ )内で文字j gramを生成

chain.from_iterableをつけない場合、map関数はiteratorを返すため、要素がmapオブジェクト(<map object at xxxxxxxxxxx>)のiteratorを返してしまう。

その場合、返り値が

[<beg gramのmap object>, <beg+1 gramのmap object>, ... , <end gramのmap object>]

となり、ngramの文字一覧としては扱いにくい。
そこで、 chain.from_iterable を利用し、生成した全ての文字ngramをiteratorで返すことができる。

# 3,4,5,6gramを生成する
list(char_ngrams("0123456789", 3, 6))
>> ['012', '123', '234', '345', '456', '567', '678', '789', '0123', '1234', '2345', '3456', '4567', '5678', '6789', '01234', '12345', '23456', '34567', '45678', '56789', '012345', '123456', '234567', '345678', '456789']

まとめ

chainはiteratorの要素を繋げるときに便利
多次元配列や組み込み関数の返り値など、箱入れになった要素を平らに変換できる

https://docs.python.org/ja/3/faq/general.html#what-is-python ↩