Source code for pyjams.pack

#!/usr/bin/env python
"""
Mimics Fortran intrinsic pack and unpack functions

This module was written by Matthias Cuntz while at Department of
Computational Hydrosystems, Helmholtz Centre for Environmental
Research - UFZ, Leipzig, Germany, and continued while at Institut
National de Recherche pour l'Agriculture, l'Alimentation et
l'Environnement (INRAE), Nancy, France.

:copyright: Copyright 2009-2023 Matthias Cuntz, see AUTHORS.rst for details.
:license: MIT License, see LICENSE for details.

.. moduleauthor:: Matthias Cuntz

The following functions are provided

.. autosummary::
   pack
   unpack

History
    * Written Jul 2009 by Matthias Cuntz (mc (at) macu (dot) de)
    * Ported to Python 3, Feb 2013, Matthias Cuntz
    * Assert matching mask and array dimensions, Apr 2014, Matthias Cuntz
    * Code refactoring, Sep 2021, Matthias Cuntz
    * Port to pyjams, May 2023, Matthias Cuntz
    * Rename value keyword to fill_value as in numpy masked arrays,
      May 2023, Matthias Cuntz
    * Assert output type equals input type, May 2023, Matthias Cuntz

"""
import numpy as np


__all__ = ['pack', 'unpack']


[docs]def pack(array, mask): """ Mimics Fortran intrinsic pack (without optional vector) Packs the last dimensions of an arbitrary shaped array into a one dimensional array under a mask. The mask can have any number of dimensions up to the array dimension. Parameters ---------- array : array ND-array to be packed mask : array Boolean ND-array with number of dimensions <= `array` dimensions Returns ------- array `array` with reduced dimensions by number of mask dimensions minus one. Last dimension has only elements that correspond to elements of `mask == True`. Notes ----- Result is undefined if all mask values are False. Examples -------- Create some data for example an island in the middle of an ocean >>> import numpy as np >>> island = np.array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], ... [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], ... [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], ... [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], ... [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]) Pack array to keep only the elements of the island >>> mask = island == 1.0 >>> pisland = pack(island, mask) >>> print(pisland.size, pisland.shape) 9 (9,) >>> print(island.sum(), pisland.sum()) 9.0 9.0 >>> print(pisland) [1. 1. 1. 1. 1. 1. 1. 1. 1.] Create data on the ocean and on the island for 2 time steps >>> tshape = (2, *island.shape) >>> temp = np.arange(np.prod(tshape)).reshape(tshape) Pack array to keep only the elements of the island >>> ptemp = pack(temp, mask) >>> print(ptemp.size, ptemp.shape) 18 (2, 9) >>> print(ptemp) [[13 14 15 23 24 25 33 34 35] [63 64 65 73 74 75 83 84 85]] """ dmask = mask.shape ndmask = np.ndim(mask) nmask = mask.size darray = array.shape ndarray = np.ndim(array) narray = array.size # Check array and mask assert ndarray >= ndmask, (f'Input array has less dimensions {ndarray}' f' then mask {ndmask}') k = 0 while k > -ndmask: k -= 1 assert dmask[k] == darray[k], ( f'Input array and mask must have the same last dimensions.' f'Array: {darray}, Mask: {dmask}' ) # Make array and mask 1d farray = array.ravel() # flat in Fortran=column-major mode fmask = mask.ravel() afmask = np.empty(narray, dtype=bool) nn = narray // nmask k = 0 while k < nn: afmask[k * nmask: (k + 1) * nmask] = fmask[:] k += 1 # Mask array and reshape afarray = farray[afmask] dout = list(darray) k = 0 while k < ndmask: del dout[-1] k += 1 nnmask = int(mask.sum()) dout.append(nnmask) out = np.reshape(afarray, dout) return out
[docs]def unpack(array, mask, fill_value=0): """ Mimics Fortran intrinsic unpack (without optional field) Unpacks the last dimension into several dimensions under a mask. The unpacked elements will be set to a user-defined fill_value. The mask can have any number of dimensions up to the array dimensions. Parameters ---------- array : array ND-array to be unpacked mask : array Boolean ND-array fill_value : float, optional Value of the new elements (default: 0) Returns ------- array `array` having more dimensions by number of mask dimensions minus one. The new elements have the user-defined `fill_value`. Notes ----- Result is undefined if all mask values are False. Examples -------- Create some data for example an island in the middle of an ocean >>> import numpy as np >>> island = np.array([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], ... [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], ... [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], ... [0., 0., 0., 1., 1., 1., 0., 0., 0., 0.], ... [0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]]) >>> print(island.size, island.shape) 50 (5, 10) Pack array to keep only the elements of the island >>> mask = island == 1.0 >>> pisland = pack(island, mask) >>> print(pisland.size, pisland.shape) 9 (9,) Unpack the ocean and island data again >>> uisland = unpack(pisland, mask) >>> print(uisland.size, uisland.shape) 50 (5, 10) >>> print(np.all(uisland == island)) True Unpack the ocean and island data, filling ocean with -9 >>> uisland = unpack(pisland, mask, -9) >>> print(uisland.size, uisland.shape) 50 (5, 10) >>> print(uisland) [[-9. -9. -9. -9. -9. -9. -9. -9. -9. -9.] [-9. -9. -9. 1. 1. 1. -9. -9. -9. -9.] [-9. -9. -9. 1. 1. 1. -9. -9. -9. -9.] [-9. -9. -9. 1. 1. 1. -9. -9. -9. -9.] [-9. -9. -9. -9. -9. -9. -9. -9. -9. -9.]] Create data on the ocean and on the island for 2 time steps >>> tshape = (2, *island.shape) >>> temp = np.arange(np.prod(tshape)).reshape(tshape) >>> print(temp.size, temp.shape) 100 (2, 5, 10) Pack array to keep only the elements of the island >>> ptemp = pack(temp, mask) >>> print(ptemp.size, ptemp.shape) 18 (2, 9) Unpack the ocean and island data. Note the ocean will have the fill_value -1 not its original data. >>> utemp = unpack(ptemp, mask, fill_value=-1.) >>> print(utemp.size, utemp.shape) 100 (2, 5, 10) >>> ii = np.where(utemp != -1.) >>> print(np.all(utemp[ii] == temp[ii])) True """ dmask = np.shape(mask) ndmask = len(dmask) nmask = mask.size darray = np.shape(array) ndarray = len(darray) narray = array.size # Check array and mask ntmask = int(mask.sum()) assert darray[-1] == ntmask, ( f'Last dimension of input array {darray[-1]} must have same size' f' as true values in mask {ntmask}.' ) # Make multi mask array array icount = nmask for i in range(ndarray - 1): icount *= darray[i] mask1d = np.ravel(mask) masknd = np.empty(icount, dtype='bool') for i in range(icount // nmask): masknd[i * nmask:(i + 1) * nmask] = mask1d[:] # Make indeces index = np.arange(icount) ii = index[masknd] if len(ii) != narray: raise ValueError(f'Index creation failed. Index {len(ii)} Array' f' {narray}.') # Flat output array array1d = np.ravel(array) arraynd = np.ones(icount, dtype=array.dtype) * fill_value arraynd[ii] = array1d[:] # Reshaped output array newdim = list(darray[0:-1]) for i in range(ndmask): newdim.append(dmask[i]) out = np.reshape(arraynd, newdim) return out
if __name__ == '__main__': import doctest doctest.testmod(optionflags=doctest.NORMALIZE_WHITESPACE)