Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment