<s> import copy <EOL> import numpy as np <EOL> import matplotlib . pyplot as plt <EOL> import matplotlib . colors <EOL> from . import plotutil <EOL> from . plotutil import bc_color_dict <EOL> class ModelMap ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , sr = None , ax = None , model = None , dis = None , layer = <NUM_LIT:0> , <EOL> extent = None , xul = None , yul = None , rotation = None ) : <EOL> self . model = model <EOL> self . layer = layer <EOL> self . dis = dis <EOL> self . sr = None <EOL> if sr is not None : <EOL> self . sr = copy . deepcopy ( sr ) <EOL> elif dis is not None : <EOL> self . sr = copy . deepcopy ( dis . sr ) <EOL> elif model is not None : <EOL> self . sr = copy . deepcopy ( model . dis . sr ) <EOL> if xul is not None : <EOL> self . sr . xul = xul <EOL> if yul is not None : <EOL> self . sr . yul = yul <EOL> if rotation is not None : <EOL> self . sr . rotation = rotation <EOL> if ax is None : <EOL> try : <EOL> self . ax = plt . gca ( ) <EOL> self . ax . set_aspect ( '<STR_LIT>' ) <EOL> except : <EOL> self . ax = plt . subplot ( <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> , aspect = '<STR_LIT>' , axisbg = "<STR_LIT>" ) <EOL> else : <EOL> self . ax = ax <EOL> if extent is not None : <EOL> self . _extent = extent <EOL> else : <EOL> self . _extent = None <EOL> return <EOL> @ property <EOL> def extent ( self ) : <EOL> if self . _extent is None : <EOL> self . _extent = self . sr . get_extent ( ) <EOL> return self . _extent <EOL> def plot_array ( self , a , masked_values = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if a . ndim == <NUM_LIT:3> : <EOL> plotarray = a [ self . layer , : , : ] <EOL> elif a . ndim == <NUM_LIT:2> : <EOL> plotarray = a <EOL> else : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> if masked_values is not None : <EOL> for mval in masked_values : <EOL> plotarray = np . ma . masked_equal ( plotarray , mval ) <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> quadmesh = ax . pcolormesh ( self . sr . xgrid , self . sr . ygrid , plotarray , <EOL> ** kwargs ) <EOL> ax . set_xlim ( self . extent [ <NUM_LIT:0> ] , self . extent [ <NUM_LIT:1> ] ) <EOL> ax . set_ylim ( self . extent [ <NUM_LIT:2> ] , self . extent [ <NUM_LIT:3> ] ) <EOL> return quadmesh <EOL> def contour_array ( self , a , masked_values = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if a . ndim == <NUM_LIT:3> : <EOL> plotarray = a [ self . layer , : , : ] <EOL> elif a . ndim == <NUM_LIT:2> : <EOL> plotarray = a <EOL> else : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> if masked_values is not None : <EOL> for mval in masked_values : <EOL> plotarray = np . ma . masked_equal ( plotarray , mval ) <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if '<STR_LIT>' in kwargs . keys ( ) : <EOL> if '<STR_LIT>' in kwargs . keys ( ) : <EOL> cmap = kwargs . pop ( '<STR_LIT>' ) <EOL> cmap = None <EOL> contour_set = ax . contour ( self . sr . xcentergrid , self . sr . ycentergrid , <EOL> plotarray , ** kwargs ) <EOL> ax . set_xlim ( self . extent [ <NUM_LIT:0> ] , self . extent [ <NUM_LIT:1> ] ) <EOL> ax . set_ylim ( self . extent [ <NUM_LIT:2> ] , self . extent [ <NUM_LIT:3> ] ) <EOL> return contour_set <EOL> def plot_inactive ( self , ibound = None , color_noflow = '<STR_LIT>' , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if ibound is None : <EOL> bas = self . model . get_package ( '<STR_LIT>' ) <EOL> ibound = bas . ibound . array <EOL> plotarray = np . zeros ( ibound . shape , dtype = np . int ) <EOL> idx1 = ( ibound == <NUM_LIT:0> ) <EOL> plotarray [ idx1 ] = <NUM_LIT:1> <EOL> plotarray = np . ma . masked_equal ( plotarray , <NUM_LIT:0> ) <EOL> cmap = matplotlib . colors . ListedColormap ( [ '<STR_LIT:0>' , color_noflow ] ) <EOL> bounds = [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> ] <EOL> norm = matplotlib . colors . BoundaryNorm ( bounds , cmap . N ) <EOL> quadmesh = self . plot_array ( plotarray , cmap = cmap , norm = norm , ** kwargs ) <EOL> return quadmesh <EOL> def plot_ibound ( self , ibound = None , color_noflow = '<STR_LIT>' , color_ch = '<STR_LIT>' , <EOL> ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if ibound is None : <EOL> bas = self . model . get_package ( '<STR_LIT>' ) <EOL> ibound = bas . ibound . array <EOL> plotarray = np . zeros ( ibound . shape , dtype = np . int ) <EOL> idx1 = ( ibound == <NUM_LIT:0> ) <EOL> idx2 = ( ibound < <NUM_LIT:0> ) <EOL> plotarray [ idx1 ] = <NUM_LIT:1> <EOL> plotarray [ idx2 ] = <NUM_LIT:2> <EOL> plotarray = np . ma . masked_equal ( plotarray , <NUM_LIT:0> ) <EOL> cmap = matplotlib . colors . ListedColormap ( [ '<STR_LIT:0>' , color_noflow , color_ch ] ) <EOL> bounds = [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:3> ] <EOL> norm = matplotlib . colors . BoundaryNorm ( bounds , cmap . N ) <EOL> quadmesh = self . plot_array ( plotarray , cmap = cmap , norm = norm , ** kwargs ) <EOL> return quadmesh <EOL> def plot_grid ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if '<STR_LIT>' not in kwargs : <EOL> kwargs [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> lc = self . get_grid_line_collection ( ** kwargs ) <EOL> ax . add_collection ( lc ) <EOL> ax . set_xlim ( self . extent [ <NUM_LIT:0> ] , self . extent [ <NUM_LIT:1> ] ) <EOL> ax . set_ylim ( self . extent [ <NUM_LIT:2> ] , self . extent [ <NUM_LIT:3> ] ) <EOL> return lc <EOL> def plot_bc ( self , ftype = None , package = None , kper = <NUM_LIT:0> , color = None , plotAll = False , <EOL> ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if package is not None : <EOL> p = package <EOL> elif self . model is not None : <EOL> if ftype is None : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> p = self . model . get_package ( ftype ) <EOL> else : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> try : <EOL> mflist = p . stress_period_data [ kper ] <EOL> except Exception as e : <EOL> raise Exception ( '<STR_LIT>' + str ( e ) ) <EOL> if mflist is None : <EOL> return None <EOL> nlay = self . model . nlay <EOL> plotarray = np . zeros ( ( nlay , self . sr . nrow , self . sr . ncol ) , dtype = np . int ) <EOL> if plotAll : <EOL> idx = [ mflist [ '<STR_LIT:i>' ] , mflist [ '<STR_LIT>' ] ] <EOL> pa = np . zeros ( ( self . sr . nrow , self . sr . ncol ) , dtype = np . int ) <EOL> pa [ idx ] = <NUM_LIT:1> <EOL> for k in range ( nlay ) : <EOL> plotarray [ k , : , : ] = pa . copy ( ) <EOL> else : <EOL> idx = [ mflist [ '<STR_LIT:k>' ] , mflist [ '<STR_LIT:i>' ] , mflist [ '<STR_LIT>' ] ] <EOL> plotarray [ idx ] = <NUM_LIT:1> <EOL> plotarray = np . ma . masked_equal ( plotarray , <NUM_LIT:0> ) <EOL> if color is None : <EOL> if ftype in bc_color_dict : <EOL> c = bc_color_dict [ ftype ] <EOL> else : <EOL> c = bc_color_dict [ '<STR_LIT:default>' ] <EOL> else : <EOL> c = color <EOL> cmap = matplotlib . colors . ListedColormap ( [ '<STR_LIT:0>' , c ] ) <EOL> bounds = [ <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:2> ] <EOL> norm = matplotlib . colors . BoundaryNorm ( bounds , cmap . N ) <EOL> quadmesh = self . plot_array ( plotarray , cmap = cmap , norm = norm , ** kwargs ) <EOL> return quadmesh <EOL> def plot_shapefile ( self , shp , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> patch_collection = plotutil . plot_shapefile ( shp , ax , ** kwargs ) <EOL> return patch_collection <EOL> def plot_discharge ( self , frf , fff , dis = None , flf = None , head = None , istep = <NUM_LIT:1> , jstep = <NUM_LIT:1> , <EOL> normalize = False , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> pivot = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> pivot = '<STR_LIT>' <EOL> if dis is None : <EOL> if self . model is not None : <EOL> dis = self . model . dis <EOL> else : <EOL> print ( "<STR_LIT>" ) <EOL> return <EOL> ib = self . model . bas6 . ibound . array <EOL> delr = dis . delr . array <EOL> delc = dis . delc . array <EOL> top = dis . top . array <EOL> botm = dis . botm . array <EOL> nlay , nrow , ncol = botm . shape <EOL> laytyp = None <EOL> hnoflo = <NUM_LIT> <EOL> hdry = <NUM_LIT> <EOL> if self . model is not None : <EOL> lpf = self . model . get_package ( '<STR_LIT>' ) <EOL> if lpf is not None : <EOL> laytyp = lpf . laytyp . array <EOL> hdry = lpf . hdry <EOL> bas = self . model . get_package ( '<STR_LIT>' ) <EOL> if bas is not None : <EOL> hnoflo = bas . hnoflo <EOL> if head is None or laytyp is None : <EOL> head = np . zeros ( botm . shape , np . float32 ) <EOL> laytyp = np . zeros ( ( nlay ) , dtype = np . int ) <EOL> sat_thk = plotutil . saturated_thickness ( head , top , botm , laytyp , <EOL> [ hnoflo , hdry ] ) <EOL> qx , qy , qz = plotutil . centered_specific_discharge ( frf , fff , flf , delr , <EOL> delc , sat_thk ) <EOL> u = qx [ self . layer , : , : ] <EOL> v = qy [ self . layer , : , : ] <EOL> x = self . sr . xcentergrid [ : : istep , : : jstep ] <EOL> y = self . sr . ycentergrid [ : : istep , : : jstep ] <EOL> u = u [ : : istep , : : jstep ] <EOL> v = v [ : : istep , : : jstep ] <EOL> if normalize : <EOL> vmag = np . sqrt ( u ** <NUM_LIT> + v ** <NUM_LIT> ) <EOL> idx = vmag > <NUM_LIT:0.> <EOL> u [ idx ] /= vmag [ idx ] <EOL> v [ idx ] /= vmag [ idx ] <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> idx = ( ib [ self . layer , : : istep , : : jstep ] == <NUM_LIT:0> ) <EOL> u [ idx ] = np . nan <EOL> v [ idx ] = np . nan <EOL> urot , vrot = self . sr . rotate ( u , v , self . sr . rotation ) <EOL> quiver = ax . quiver ( x , y , urot , vrot , pivot = pivot , ** kwargs ) <EOL> return quiver <EOL> def plot_pathline ( self , pl , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> from matplotlib . collections import LineCollection <EOL> if isinstance ( pl , np . ndarray ) : <EOL> pl = [ pl ] <EOL> if '<STR_LIT>' in kwargs : <EOL> kon = kwargs . pop ( '<STR_LIT>' ) <EOL> if isinstance ( kon , str ) : <EOL> if kon . lower ( ) == '<STR_LIT:all>' : <EOL> kon = - <NUM_LIT:1> <EOL> else : <EOL> kon = self . layer <EOL> else : <EOL> kon = self . layer <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if '<STR_LIT>' not in kwargs : <EOL> kwargs [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> linecol = [ ] <EOL> for p in pl : <EOL> vlc = [ ] <EOL> x0r , y0r = self . sr . rotate ( p [ '<STR_LIT:x>' ] , p [ '<STR_LIT:y>' ] , self . sr . rotation , <NUM_LIT:0.> , self . sr . yedge [ <NUM_LIT:0> ] ) <EOL> x0r += self . sr . xul <EOL> y0r += self . sr . yul - self . sr . yedge [ <NUM_LIT:0> ] <EOL> arr = np . vstack ( ( x0r , y0r ) ) . T <EOL> if kon >= <NUM_LIT:0> : <EOL> arr = np . ma . masked_where ( ( p [ '<STR_LIT:k>' ] != kon ) , arr ) <EOL> if not arr . mask . all ( ) : <EOL> linecol . append ( arr ) <EOL> lc = None <EOL> if len ( linecol ) > <NUM_LIT:0> : <EOL> lc = LineCollection ( linecol , ** kwargs ) <EOL> ax . add_collection ( lc ) <EOL> return lc <EOL> def plot_endpoint ( self , ep , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> if '<STR_LIT>' in kwargs : <EOL> kon = kwargs . pop ( '<STR_LIT>' ) <EOL> if isinstance ( kon , str ) : <EOL> if kon . lower ( ) == '<STR_LIT:all>' : <EOL> kon = - <NUM_LIT:1> <EOL> else : <EOL> kon = self . layer <EOL> else : <EOL> kon = self . layer <EOL> if '<STR_LIT>' in kwargs : <EOL> ax = kwargs . pop ( '<STR_LIT>' ) <EOL> else : <EOL> ax = self . ax <EOL> if '<STR_LIT:c>' not in kwargs : <EOL> c = ep [ '<STR_LIT:time>' ] <EOL> else : <EOL> c = np . empty ( ( ep . shape [ <NUM_LIT:0> ] ) , dtype = "<STR_LIT>" ) <EOL> c . fill ( kwargs . pop ( '<STR_LIT:c>' ) ) <EOL> if '<STR_LIT:s>' not in kwargs : <EOL> s = <NUM_LIT> <EOL> else : <EOL> s = float ( kwargs . pop ( '<STR_LIT:s>' ) ) ** <NUM_LIT> <EOL> createcb = False <EOL> if '<STR_LIT>' in kwargs : <EOL> createcb = kwargs . pop ( '<STR_LIT>' ) <EOL> colorbar_label = '<STR_LIT>' <EOL> if '<STR_LIT>' in kwargs : <EOL> colorbar_label = kwargs . pop ( '<STR_LIT>' ) <EOL> shrink = <NUM_LIT:1.> <EOL> if '<STR_LIT>' in kwargs : <EOL> shrink = float ( kwargs . pop ( '<STR_LIT>' ) ) <EOL> x0r , y0r = self . sr . rotate ( ep [ '<STR_LIT:x>' ] , ep [ '<STR_LIT:y>' ] , self . sr . rotation , <NUM_LIT:0.> , self . sr . yedge [ <NUM_LIT:0> ] ) <EOL> x0r += self . sr . xul <EOL> y0r += self . sr . yul - self . sr . yedge [ <NUM_LIT:0> ] <EOL> arr = np . vstack ( ( x0r , y0r ) ) . T <EOL> if kon >= <NUM_LIT:0> : <EOL> c = np . ma . masked_where ( ( ep [ '<STR_LIT:k>' ] != kon ) , c ) <EOL> sp = plt . scatter ( arr [ : , <NUM_LIT:0> ] , arr [ : , <NUM_LIT:1> ] , c = c , s = s , ** kwargs ) <EOL> if createcb : <EOL> cb = plt . colorbar ( sp , shrink = shrink ) <EOL> cb . set_label ( colorbar_label ) <EOL> return sp <EOL> def get_grid_line_collection ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> from matplotlib . collections import LineCollection <EOL> lc = LineCollection ( self . sr . get_grid_lines ( ) , ** kwargs ) <EOL> return lc </s>
