#NEWLINE# def saturated_thickness( head , top , botm , laytyp , mask_values = None ) : #NEWLINE# #INDENT# '\n    Calculate the saturated thickness.\n\n    Parameters\n    ----------\n    head : numpy.ndarray\n        head array\n    top : numpy.ndarray\n        top array of shape (nrow, ncol)\n    botm : numpy.ndarray\n        botm array of shape (nlay, nrow, ncol)\n    laytyp : numpy.ndarray\n        confined (0) or convertible (1) of shape (nlay)\n    mask_values : list of floats\n        If head is one of these values, then set sat to top - bot\n\n    Returns\n    -------\n    sat_thk : numpy.ndarray\n        Saturated thickness of shape (nlay, nrow, ncol).\n\n    ' #NEWLINE# ( nlay , nrow , ncol ) = head . shape #NEWLINE# sat_thk = np . empty ( head . shape , dtype = head . dtype ) #NEWLINE# for k in range ( nlay ) : #NEWLINE# #INDENT# if ( k == 0 ) : #NEWLINE# #INDENT# t = top #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# t = botm [ ( k - 1 ) , : , : ] #NEWLINE# #UNINDENT# sat_thk [ k , : , : ] = ( t - botm [ k , : , : ] ) #NEWLINE# #UNINDENT# for k in range ( nlay ) : #NEWLINE# #INDENT# if ( laytyp [ k ] != 0 ) : #NEWLINE# #INDENT# dh = np . zeros ( ( nrow , ncol ) , dtype = head . dtype ) #NEWLINE# s = sat_thk [ k , : , : ] #NEWLINE# for mv in mask_values : #NEWLINE# #INDENT# idx = ( head [ k , : , : ] == top ) #NEWLINE# dh [ idx ] = s [ idx ] #NEWLINE# #UNINDENT# if ( k == 0 ) : #NEWLINE# #INDENT# t = top #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# t = botm [ ( k - 1 ) , : , : ] #NEWLINE# #UNINDENT# t = np . where ( ( head [ k , : , : ] > t ) , t , head [ k , : , : ] ) #NEWLINE# dh = np . where ( ( dh == 0 ) , ( t - botm [ k , : , : ] ) , dh ) #NEWLINE# sat_thk [ k , : , : ] = dh [ : , : ] #NEWLINE# #UNINDENT# #UNINDENT# return sat_thk