xx_calc_modified_v2.py

import numpy


def get_masked_arr(arr, fill_val):
    '''
    If a masked array is passed, this function does nothing.
    If a filled array is passed (fill_value must be passed also), it will be transformed into a masked array.
    
    '''
    if isinstance(arr, numpy.ma.MaskedArray):               # numpy.ma.MaskedArray
        masked_arr = arr
    else:                                                   # numpy.ndarray
        if (fill_val==None):
            raise(ValueError('If input array is not a masked array, a "fill_value" must be provided.'))
        mask_arr = (arr==fill_val)
        masked_arr = numpy.ma.masked_array(arr, mask=mask_arr, fill_value=fill_val)
    
    return masked_arr


def SU_calculation(arr, fill_val=None):
    '''
    Calculates the SU indice: summer days (daily maximum temperature > 25 degrees Celsius) [days].
    
    :param arr: daily maximum temperature (e.g. "tasmax") in Kelvin
    :type arr: numpy.ndarray (3D) or numpy.ma.MaskedArray (3D)
    :param fill_val: fill value 
    :type fill_val: float
    
    :rtype: numpy.ndarray (2D)        (if "arr" is numpy.ndarray)
         or numpy.ma.MaskedArray (2D) (if "arr" is numpy.ma.MaskedArray)
    '''
    
    #t = 25          # temperature threshold (degree Celsius)
    #T = t + 273.15  # Celsius -> Kelvin
    #
    
    t=10
    
    data = get_masked_arr(arr, fill_val) # numpy.ma.MaskedArray with fill_value=fill_val (if numpy.ndarray passed) or fill_value=arr.fill_value (if numpy.ma.MaskedArray is passed)
    
    data = (data>t)                     # fill_value not changed
    SU = data.sum(axis=0)               # fill_value is changed: SU is a new numpy.ma.MaskedArray with default fill_value=999999 (!) => next line is to keep the fill_value of data 
    numpy.ma.set_fill_value(SU, data.fill_value) 
    
    if not isinstance(arr, numpy.ma.MaskedArray):
        SU = SU.filled(fill_value=data.fill_value) # numpy.ndarray filled with input fill_val
    
    return SU



#################### TEST ####################################


def test_masked_array():
    arr = numpy.array([   [ [1, 5, 15, 9999],[4, 11, 2, 10], [9999, 0, 19, 5] ], [ [0, 16, 1, 9999],[9999, 11, 5, 22], [3, 0, 7, 5] ]    ]) # 2 x 3 x 4
    
    mask_arr = (arr==9999)
    
    arr_masked = numpy.ma.masked_array(arr, mask=mask_arr, fill_value=4444)

    res = SU_calculation(arr_masked, fill_val=9999)

    assert(type(res) == numpy.ma.MaskedArray)
    to_test = numpy.ma.array([[0,1,1,0],[0,2,0,1],[0,0,1,0]],mask=False)
    to_test.mask[0,-1] = True
    assert((res.data == to_test.data).all())
    assert((res.mask == to_test.mask).all())
    assert(res.fill_value == 9999)
    
def test_normal_filled_array():
    
    arr = numpy.array([[[1,6,99,18],[56,7,2,99]]])
    
    res = SU_calculation(arr,fill_val=99)
    
    assert((res == numpy.array([[0,0,99,1],[1,0,0,99]])).all())