a ߙfb>@sTdZddlZddlmZgdZddZdd Zeejfd d Z edd dZ dS)z= This module includes helper functions for array operations. N)support_nddata)reshape_as_blocks block_reduceblock_replicatecCst|}t|}t|dkr*td|jdkrNt|dkrNt||j}t||jkrdtd|t }t||krtd||fS)Nrz-block_size elements must be strictly positiverzTblock_size must be a scalar or have the same length as the number of data dimensionsz$block_size elements must be integers) npZ asanyarrayZ atleast_1dany ValueErrorndimlenrepeatZastypeint)data block_sizeZblock_size_intr4lib/python3.9/site-packages/astropy/nddata/blocks.py_process_block_inputs s   rcCst||\}}tt|j|dkr.tdt|j|}tddt||D}tt dt |d}tt dt |d}| | ||S)a Reshape a data array into blocks. This is useful to efficiently apply functions on block subsets of the data instead of using loops. The reshaped array is a view of the input data array. .. versionadded:: 4.1 Parameters ---------- data : ndarray The input data array. block_size : int or array-like (int) The integer block size along each axis. If ``block_size`` is a scalar and ``data`` has more than one dimension, then ``block_size`` will be used for for every axis. Each dimension of ``block_size`` must divide evenly into the corresponding dimension of ``data``. Returns ------- output : ndarray The reshaped array as a view of the input ``data`` array. Examples -------- >>> import numpy as np >>> from astropy.nddata import reshape_as_blocks >>> data = np.arange(16).reshape(4, 4) >>> data array([[ 0, 1, 2, 3], [ 4, 5, 6, 7], [ 8, 9, 10, 11], [12, 13, 14, 15]]) >>> reshape_as_blocks(data, (2, 2)) array([[[[ 0, 1], [ 4, 5]], [[ 2, 3], [ 6, 7]]], [[[ 8, 9], [12, 13]], [[10, 11], [14, 15]]]]) rzXEach dimension of block_size must divide evenly into the corresponding dimension of datacss|]}|D] }|Vq qdS)Nr).0Zijkrrr Yz$reshape_as_blocks..r) rrrmodshaper arraytuplezipranger ZreshapeZ transpose)rrnblocksZ new_shapeZ nblocks_idxZ block_idxrrrr"s0rcCst||\}}t|j|}||}t|jD]>}|j|||kr0|d|}|d||}|d|}q0t||}tt|j|j}|||dS)a Downsample a data array by applying a function to local blocks. If ``data`` is not perfectly divisible by ``block_size`` along a given axis then the data will be trimmed (from the end) along that axis. Parameters ---------- data : array-like The data to be resampled. block_size : int or array-like (int) The integer block size along each axis. If ``block_size`` is a scalar and ``data`` has more than one dimension, then ``block_size`` will be used for for every axis. func : callable, optional The method to use to downsample the data. Must be a callable that takes in a `~numpy.ndarray` along with an ``axis`` keyword, which defines the axis or axes along which the function is applied. The ``axis`` keyword must accept multiple axes as a tuple. The default is `~numpy.sum`, which provides block summation (and conserves the data sum). Returns ------- output : array-like The resampled data. Examples -------- >>> import numpy as np >>> from astropy.nddata import block_reduce >>> data = np.arange(16).reshape(4, 4) >>> block_reduce(data, 2) # doctest: +FLOAT_CMP array([[10, 18], [42, 50]]) >>> block_reduce(data, 2, func=np.mean) # doctest: +FLOAT_CMP array([[ 2.5, 4.5], [ 10.5, 12.5]]) rNaxis) rrrrrr Zswapaxesrr)rrfuncrZ size_initr Zreshapedrrrr`s.  rTcCsLt||\}}t|jD]}tj||||d}q|rH|tt|}|S)a Upsample a data array by block replication. Parameters ---------- data : array-like The data to be block replicated. block_size : int or array-like (int) The integer block size along each axis. If ``block_size`` is a scalar and ``data`` has more than one dimension, then ``block_size`` will be used for for every axis. conserve_sum : bool, optional If `True` (the default) then the sum of the output block-replicated data will equal the sum of the input ``data``. Returns ------- output : array-like The block-replicated data. Examples -------- >>> import numpy as np >>> from astropy.nddata import block_replicate >>> data = np.array([[0., 1.], [2., 3.]]) >>> block_replicate(data, 2) # doctest: +FLOAT_CMP array([[0. , 0. , 0.25, 0.25], [0. , 0. , 0.25, 0.25], [0.5 , 0.5 , 0.75, 0.75], [0.5 , 0.5 , 0.75, 0.75]]) >>> block_replicate(data, 2, conserve_sum=False) # doctest: +FLOAT_CMP array([[0., 0., 1., 1.], [0., 0., 1., 1.], [2., 2., 3., 3.], [2., 2., 3., 3.]]) r)rrr rr floatZprod)rrZ conserve_sumirrrrs *r)T) __doc__ZnumpyrZ decoratorsr__all__rrsumrrrrrrs >>