ó ¦–Õ\c@s6ddlZddlZyddlZWnek rGddlZnXddlmZmZddl m Z m Z ddl m Z ddlmZddlmZd„Zd „Zd d „Zd „Zd e fd„ƒYZd„Zd„Zdd„Zdd„Zd„Zd„Zddfdd„ZdS(iÿÿÿÿNi(tcoretutils(tapplytMapping(t to_task_dask(tHighLevelGraph(tSubgraphCallablecs›t|tƒr)‡fd†|jƒDƒSt|ƒtttfkrpt|ƒg|D]}t|ˆƒ^qTƒSy ˆ|SWntt fk r–|SXdS(sË Create a new task with the values substituted This is like dask.core.subs, but takes a dict of many substitutions to perform simultaneously. It is not as concerned with micro performance. cs(i|]\}}t|ˆƒ|“qS((tsubs(t.0tktv(t substitution(s-lib/python2.7/site-packages/dask/blockwise.pys s N( t isinstancetdicttitemsttypettupletlisttsetRtKeyErrort TypeError(ttaskR tx((R s-lib/python2.7/site-packages/dask/blockwise.pyRs, cCs=|dkr|Stg|D]}|j||ƒ^qƒSdS(s> A simple subs function that works both on tuples and strings N(tNoneRtget(tindR tc((s-lib/python2.7/site-packages/dask/blockwise.pyt index_subs s t_cCs |d|S(Ns%d((titprefix((s-lib/python2.7/site-packages/dask/blockwise.pytblockwise_token(sc s;|jdƒ}|jddƒ}|jdiƒ}|jdgƒ}t|ƒ}d„|ddd…Dƒt|ƒB} d „tt| ƒƒDƒ‰tt|ƒˆƒ}g|ddd…D]$} | dk rÛt| ƒn| ^q½|ddd…@s  iNicSs(i|]\}}t|dƒ|“qS(t.(R(RRR ((s-lib/python2.7/site-packages/dask/blockwise.pys Cs cs/i|]%\}}|t|fˆƒd“qS(i(R(RR R (tsub(s-lib/python2.7/site-packages/dask/blockwise.pys Js cSs)h|]}|jƒD] }|’qqS((t__dask_layers__(RRtn((s-lib/python2.7/site-packages/dask/blockwise.pys Ts css|]}t|ƒVqdS(N(R(RR((s-lib/python2.7/site-packages/dask/blockwise.pys Ws(N(tpopRRRt enumeratetsortedRRRttoolzt partitiontrangetlenR tzipRtmapRRt Blockwise(tfunctoutputtoutput_indicest arrind_pairstkwargsR R!R"R#tunique_indicestatargpairstnameRtinputstindextinputs_indicestnew_keyst new_tokensR R tindicestkeystsubgrapht_keystkwargs2((R&s-lib/python2.7/site-packages/dask/blockwise.pyt blockwise,sJ   H9%%5.! !R2cBsPeZdZddd„Zed„ƒZd„Zd„Zd„Z d„Z RS(s£ Tensor Operation This is a lazily constructed mapping for tensor operation graphs. This defines a dictionary using an operation and an indexing pattern. It is built for many operations like elementwise, transpose, tensordot, and so on. We choose to keep these as symbolic mappings rather than raw dictionaries because we are able to fuse them during optimization, sometimes resulting in much lower overhead. Parameters ---------- output: str The name of the output collection. Used in keynames output_indices: tuple The output indices, like ``('i', 'j', 'k')`` used to determine the structure of the block computations dsk: dict A small graph to apply per-output-block. May include keys from the input indices. indices: Tuple[str, Tuple[str, str]] An ordered mapping from input key name, like ``'x'`` to input indices, like ``('i', 'j')`` Or includes literals, which have ``None`` for an index value numblocks: Dict[key, Sequence[int]] Number of blocks along each dimension for each input concatenate: boolean Whether or not to pass contracted dimensions as a list of inputs or a single input to the block function new_axes: Dict New index dimensions that may have been created, and their extent See Also -------- dask.blockwise.blockwise dask.array.blockwise cCs_||_t|ƒ|_||_td„|Dƒƒ|_||_||_|pUi|_dS(Ncss9|]/\}}||dk r*t|ƒn|fVqdS(N(RR(RR;R((s-lib/python2.7/site-packages/dask/blockwise.pys šs(R4RR5tdskRAR R!R"(tselfR4R5RGRAR R!R"((s-lib/python2.7/site-packages/dask/blockwise.pyt__init__•s     c Cs¤t|dƒr|jSttttt|jƒƒƒƒ}t|j |j |ƒ}t ||j |j d|j d|jd|jttj|jƒƒŒ|_|jS(Nt _cached_dictR"R R!(thasattrRJRR1RR.R/RARRGR4tmake_blockwise_graphR5R"R R!RR,tconcat(RHRBR3((s-lib/python2.7/site-packages/dask/blockwise.pyt_dict s$   $cCs |j|S(N(RN(RHtkey((s-lib/python2.7/site-packages/dask/blockwise.pyt __getitem__²scCs t|jƒS(N(titerRN(RH((s-lib/python2.7/site-packages/dask/blockwise.pyt__iter__µscCs%ttjt|jƒjƒƒƒƒS(N(tinttnptprodRt_out_numblockstvalues(RH((s-lib/python2.7/site-packages/dask/blockwise.pyt__len__¸scs–i}d„ˆjDƒ}x`ˆjjƒD]O\}}x@t|||ƒD]+\}}t|j|dƒ|ƒ||½s ics.i|]$\}}|ˆjkr||“qS((R5(RR R (RH(s-lib/python2.7/site-packages/dask/blockwise.pys Âs (RAR RR0tmaxR(RHtdRAR R R9tb((RHs-lib/python2.7/site-packages/dask/blockwise.pyRV»s  'N( t__name__t __module__t__doc__RRItpropertyRNRPRRRXRV(((s-lib/python2.7/site-packages/dask/blockwise.pyR2ps$    c#s||jdƒ}|jdd ƒ}|jdiƒ}ttjd|ƒƒ}|tkrjddlm}nt|ƒd„|DƒksŒt ‚d„|Dƒ} | t|ƒ} t ||ƒ‰x?|j ƒD]1\} } t | t ƒrït| ƒnd ˆ| >> inc = lambda x: x + 1 >>> make_blockwise_graph(inc, 'z', 'ij', 'x', 'ij', numblocks={'x': (2, 2)}) # doctest: +SKIP {('z', 0, 0): (inc, ('x', 0, 0)), ('z', 0, 1): (inc, ('x', 0, 1)), ('z', 1, 0): (inc, ('x', 1, 0)), ('z', 1, 1): (inc, ('x', 1, 1))} Simple operation on two datasets >>> add = lambda x, y: x + y >>> make_blockwise_graph(add, 'z', 'ij', 'x', 'ij', 'y', 'ij', numblocks={'x': (2, 2), ... 'y': (2, 2)}) # doctest: +SKIP {('z', 0, 0): (add, ('x', 0, 0), ('y', 0, 0)), ('z', 0, 1): (add, ('x', 0, 1), ('y', 0, 1)), ('z', 1, 0): (add, ('x', 1, 0), ('y', 1, 0)), ('z', 1, 1): (add, ('x', 1, 1), ('y', 1, 1))} Operation that flips one of the datasets >>> addT = lambda x, y: x + y.T # Transpose each chunk >>> # z_ij ~ x_ij y_ji >>> # .. .. .. notice swap >>> make_blockwise_graph(addT, 'z', 'ij', 'x', 'ij', 'y', 'ji', numblocks={'x': (2, 2), ... 'y': (2, 2)}) # doctest: +SKIP {('z', 0, 0): (add, ('x', 0, 0), ('y', 0, 0)), ('z', 0, 1): (add, ('x', 0, 1), ('y', 1, 0)), ('z', 1, 0): (add, ('x', 1, 0), ('y', 0, 1)), ('z', 1, 1): (add, ('x', 1, 1), ('y', 1, 1))} Dot product with contraction over ``j`` index. Yields list arguments >>> make_blockwise_graph(dotmany, 'z', 'ik', 'x', 'ij', 'y', 'jk', numblocks={'x': (2, 2), ... 'y': (2, 2)}) # doctest: +SKIP {('z', 0, 0): (dotmany, [('x', 0, 0), ('x', 0, 1)], [('y', 0, 0), ('y', 1, 0)]), ('z', 0, 1): (dotmany, [('x', 0, 0), ('x', 0, 1)], [('y', 0, 1), ('y', 1, 1)]), ('z', 1, 0): (dotmany, [('x', 1, 0), ('x', 1, 1)], [('y', 0, 0), ('y', 1, 0)]), ('z', 1, 1): (dotmany, [('x', 1, 0), ('x', 1, 1)], [('y', 0, 1), ('y', 1, 1)])} Pass ``concatenate=True`` to concatenate arrays ahead of time >>> make_blockwise_graph(f, 'z', 'i', 'x', 'ij', 'y', 'ij', concatenate=True, ... numblocks={'x': (2, 2), 'y': (2, 2,)}) # doctest: +SKIP {('z', 0): (f, (concatenate_axes, [('x', 0, 0), ('x', 0, 1)], (1,)), (concatenate_axes, [('y', 0, 0), ('y', 0, 1)], (1,))) ('z', 1): (f, (concatenate_axes, [('x', 1, 0), ('x', 1, 1)], (1,)), (concatenate_axes, [('y', 1, 0), ('y', 1, 1)], (1,)))} Supports Broadcasting rules >>> make_blockwise_graph(add, 'z', 'ij', 'x', 'ij', 'y', 'ij', numblocks={'x': (1, 2), ... 'y': (2, 2)}) # doctest: +SKIP {('z', 0, 0): (add, ('x', 0, 0), ('y', 0, 0)), ('z', 0, 1): (add, ('x', 0, 1), ('y', 0, 1)), ('z', 1, 0): (add, ('x', 0, 0), ('y', 1, 0)), ('z', 1, 1): (add, ('x', 0, 1), ('y', 1, 1))} Support keyword arguments with apply >>> def f(a, b=0): return a + b >>> make_blockwise_graph(f, 'z', 'i', 'x', 'i', numblocks={'x': (2,)}, b=10) # doctest: +SKIP {('z', 0): (apply, f, [('x', 0)], {'b': 10}), ('z', 1): (apply, f, [('x', 1)], {'b': 10})} Include literals by indexing with ``None`` >>> make_blockwise_graph(add, 'z', 'i', 'x', 'i', 100, None, numblocks={'x': (2,)}) # doctest: +SKIP {('z', 0): (add, ('x', 0), 100), ('z', 1): (add, ('x', 1), 100)} See Also -------- dask.array.blockwise dask.blockwise.blockwise R R!R"iiÿÿÿÿ(tconcatenate_axescSs(h|]\}}|dk r|’qS(N(R(RR;R((s-lib/python2.7/site-packages/dask/blockwise.pys 7s cSs/h|]%\}}|r|D] }|’qqS(((RRRR((s-lib/python2.7/site-packages/dask/blockwise.pys 9s ic3s+|]!}|ttˆ|ƒƒfVqdS(N(RR.(RR(tdims(s-lib/python2.7/site-packages/dask/blockwise.pys Gscss|]}|dkVqdS(iN((Rtnb((s-lib/python2.7/site-packages/dask/blockwise.pys TsN(R)RRR,R-tTruetdask.array.coreR`RtAssertionErrortbroadcast_dimensionsRR RR/t itertoolstproductR.R R0tappendt lol_tuplestanytzero_broadcast_dimensionsR*RtupdateRt ensure_dictR(#R3R4t out_indicesR6R7R R!R"R:t all_indicest dummy_indicesR R RtkeytupsttuptkeydictstdummiesRGtvaltupstkdtargstargRttupsttups2R(taxestvttktRBRtdsk2REtvals((Ras-lib/python2.7/site-packages/dask/blockwise.pyRLÅsXj ")2(  1"  ( cCs}|s |S|d|kr@t|||df|d||ƒSg||dD]&}t||f|d||ƒ^qOSdS(sD List of list of tuple keys Parameters ---------- head : tuple The known tuple so far ind : Iterable An iterable of indices not yet covered values : dict Known values for non-dummy indices dummies : dict Ranges of values for dummy indices Examples -------- >>> lol_tuples(('x',), 'ij', {'i': 1, 'j': 0}, {}) ('x', 1, 0) >>> lol_tuples(('x',), 'ij', {'i': 1}, {'j': range(3)}) [('x', 1, 0), ('x', 1, 1), ('x', 1, 2)] >>> lol_tuples(('x',), 'ij', {'i': 1}, {'j': range(3)}) [('x', 1, 0), ('x', 1, 1), ('x', 1, 2)] >>> lol_tuples(('x',), 'ijk', {'i': 1}, {'j': [0, 1, 2], 'k': [0, 1]}) # doctest: +NORMALIZE_WHITESPACE [[('x', 1, 0, 0), ('x', 1, 0, 1)], [('x', 1, 1, 0), ('x', 1, 1, 1)], [('x', 1, 2, 0), ('x', 1, 2, 1)]] iiN(Rj(theadRRWRuR ((s-lib/python2.7/site-packages/dask/blockwise.pyRjus &cCsGt|d|ƒ}x.|j|jkrB|}t|d|ƒ}qW|S(sà High level optimization of stacked Blockwise layers For operations that have multiple Blockwise operations one after the other, like ``x.T + 123`` we can fuse these into a single Blockwise operation. This happens before any actual tasks are generated, and so can reduce overhead. This finds groups of Blockwise operations that can be safely fused, and then passes them to ``rewrite_blockwise`` for rewriting. Parameters ---------- full_graph: HighLevelGraph keys: Iterable The keys of all outputs of all collections. Used to make sure that we don't fuse a layer needed by an output Returns ------- HighLevelGraph See Also -------- rewrite_blockwise RB(t_optimize_blockwiseR#(tgraphRBtout((s-lib/python2.7/site-packages/dask/blockwise.pytoptimize_blockwisežs csöd„|Dƒ}|j}tj|jƒ‰‡fd†|jDƒ}t|ƒ}i}i}tƒ}x|rè|jƒ} | |ksh| |kr˜qhn|j| ƒt|| t ƒrŸ| h} t| ƒ} xˆ| rW| jƒ‰ˆ|kr|j ˆƒqÐnt|ˆt ƒs'|j ˆƒqÐnˆ| krRˆ|krR|j ˆƒqÐn|ˆj || j kr|j ˆƒqÐnt ‡fd†|| j Dƒƒdkr»|j ˆƒqÐn| jˆƒx‰|jjˆdƒD]r} t|ˆjƒ} d„|ˆj Dƒ}tˆ| ƒdkrC| j|ƒrC| j| ƒqÞ|j | ƒqÞWqÐWtg| D]}||^qbƒ}||| ¿s cs%h|]}ˆj|ƒs|’qS((R(RR (t dependents(s-lib/python2.7/site-packages/dask/blockwise.pys Âs c3s-|]#\}}|dk r|ˆkVqdS(N(R(RR R(tdep(s-lib/python2.7/site-packages/dask/blockwise.pys âsicSs/h|]%\}}|r|D] }|’qqS(((RRRR((s-lib/python2.7/site-packages/dask/blockwise.pys ís cSs(h|]\}}|dk r|’qS(N(R(RR R ((s-lib/python2.7/site-packages/dask/blockwise.pys ÷s (((tdictsRt reverse_dictR#RRR)taddR R2RiR!tsumRARR5R/t issupersettrewrite_blockwisetextendR(t full_graphRBtkeeptlayerstrootststackR„R#tseentlayertblockwise_layerstdepsRZR5t input_indicestlt new_layer((R‡R†s-lib/python2.7/site-packages/dask/blockwise.pyR‚¾s^               )  %#  c s±d„ˆDƒ‰‡fd†ˆjƒDƒ}tj|ƒ}d„tjƒDƒ}g|jƒD]\}}|s^|^q^\}tˆ|jƒ}ˆ|j}ˆ|j } t ˆ|j ƒ} t } xR| rt } x?t|ƒD]1\‰\‰} | dkrqánˆˆkrqánt } ‡‡fd†| jƒDƒ} |jˆƒ\} }d„tˆdt|ƒdƒDƒ‰t| ˆƒ} ˆˆj}t tˆˆj|ƒƒ‰‡‡fd†|Dƒ}t t||ƒƒ}ˆj|ƒg|D]!\}}|t|ˆƒf^qò}x/ˆˆjjƒD]\}}||ˆ|s cs/i|]%}‡fd†|jDƒ|j“qS(cs4h|]*\}}|dk r|ˆkr|’qS(N(R(RRZR (R<(s-lib/python2.7/site-packages/dask/blockwise.pys s (RAR4(RR›(R<(s-lib/python2.7/site-packages/dask/blockwise.pys s css8|].}dD]!}||r(t|ƒndVq qdS(tABCDEFGHIJKLMNOPQRSTUVWXYZtN(tstr(RRZR((s-lib/python2.7/site-packages/dask/blockwise.pys scs5i|]+\}}t|iˆtˆƒ6ƒ|“qS((RR(RR R (R‡R(s-lib/python2.7/site-packages/dask/blockwise.pys 2s cSs)i|]}t|dƒt|ƒ“qS(i(R(RR((s-lib/python2.7/site-packages/dask/blockwise.pys 7s icsHh|]>\}}|dk r|D]}|ˆˆjkr"|’q"qS(N(RR5(RRtjR(R‡R<(s-lib/python2.7/site-packages/dask/blockwise.pys >s   cSs+i|]!\}}t|ƒt|ƒ“qS((R(RR R ((s-lib/python2.7/site-packages/dask/blockwise.pys ms cs(i|]\}}t|ˆƒ|“qS((R(RR R (R&(s-lib/python2.7/site-packages/dask/blockwise.pys ns cSs(h|]\}}|dk r|’qS(N(R(RR R ((s-lib/python2.7/site-packages/dask/blockwise.pys ps cs7i|]-\}}|dks*|ˆkr||“qS(N(R(RR R (t indices_check(s-lib/python2.7/site-packages/dask/blockwise.pys rs R R"R!N(#RWRR‰RgtcountRRRAR"R!R RGRctFalseR*RR)R.R/RR0R5RmRt ValueErrorRRR=RiR RR,tmergeR R2(R<R#R†tnew_index_iterR R trootRAR"R!RGtchangedRRtcurrent_dep_indicest new_indicest contractedtextraRRŸR=tcontainstnew_dskR9R[R”R›R R„((R‡RR R<R&s-lib/python2.7/site-packages/dask/blockwise.pyRs€ .     '  .  "=+cs‡fd†}tj||ƒS(s  >>> lol = [('x', 1, 0), ('x', 1, 1), ('x', 1, 2)] >>> nblocks = (4, 1, 2) # note singleton dimension in second place >>> lol = [[('x', 1, 0, 0), ('x', 1, 0, 1)], ... [('x', 1, 1, 0), ('x', 1, 1, 1)], ... [('x', 1, 2, 0), ('x', 1, 2, 1)]] >>> zero_broadcast_dimensions(lol, nblocks) # doctest: +NORMALIZE_WHITESPACE [[('x', 1, 0, 0), ('x', 1, 0, 1)], [('x', 1, 0, 0), ('x', 1, 0, 1)], [('x', 1, 0, 0), ('x', 1, 0, 1)]] See Also -------- lol_tuples cs,|dftd„t|dˆƒDƒƒS(Nicss-|]#\}}|dkr!dn|VqdS(iiN((RRRZ((s-lib/python2.7/site-packages/dask/blockwise.pys si(RR0(tt(tnblocks(s-lib/python2.7/site-packages/dask/blockwise.pytR(Rthomogeneous_deepmap(tlolR¯tf((R¯s-lib/python2.7/site-packages/dask/blockwise.pyRl{sc sIg|D]$\}}|dk r||f^q}tjgtjtj|tj|jƒƒD]'\\}}\}} t|| ƒ^q\ƒ} tjd| ƒ} td„| jƒDƒƒ} t‡fd†| jƒDƒƒ} |ròtj || ƒS| r6t t t | j ƒƒƒt dgƒk r6td| ƒ‚ntj tj| ƒS(sà Find block dimensions from arguments Parameters ---------- argpairs: iterable name, ijk index pairs numblocks: dict maps {name: number of blocks} sentinels: iterable (optional) values for singleton dimensions consolidate: func (optional) use this to reduce each set of common blocks into a smaller set Examples -------- >>> argpairs = [('x', 'ij'), ('y', 'ji')] >>> numblocks = {'x': (2, 3), 'y': (3, 2)} >>> broadcast_dimensions(argpairs, numblocks) {'i': 2, 'j': 3} Supports numpy broadcasting rules >>> argpairs = [('x', 'ij'), ('y', 'ij')] >>> numblocks = {'x': (2, 1), 'y': (1, 3)} >>> broadcast_dimensions(argpairs, numblocks) {'i': 2, 'j': 3} Works in other contexts too >>> argpairs = [('x', 'ij'), ('y', 'ij')] >>> d = {'x': ('Hello', 1), 'y': (1, (2, 3))} >>> broadcast_dimensions(argpairs, d) {'i': 'Hello', 'j': (2, 3)} icss@|]6\}}|tg|D]\}}|^qƒfVqdS(N(R(RR R RRZ((s-lib/python2.7/site-packages/dask/blockwise.pys »sc3sC|]9\}}|t|ƒdkr4|tˆƒn|fVqdS(iN(R/R(RR R (t sentinels(s-lib/python2.7/site-packages/dask/blockwise.pys ½sisShapes do not align %sN(RR,RMtjointfirstRR0tgroupbyR tvalmapRR1R/RWR£( R:R R´t consolidateR9Rt argpairs2RtindsRatLtgtg2((R´s-lib/python2.7/site-packages/dask/blockwise.pyRf‘s%1 R"1(((i(RgtnumpyRTtcytoolzR,t ImportErrorRRRt compatibilityRRtdelayedRthighlevelgraphRt optimizationRRRRRFR2RLRjR…R‚RRlRRf(((s-lib/python2.7/site-packages/dask/blockwise.pyts.       DU ° ) B {