Visualization and Plotting

class PlotMapView:
    """Map view plotting for MODFLOW models"""
    def __init__(
        self,
        model: object = None,
        ax: object = None,
        layer: int = 0,
        extent: list[float] = None,
        modelgrid: object = None,
        **kwargs
    ): ...
    
    def plot_grid(
        self,
        **kwargs
    ) -> object:
        """Plot model grid"""
        ...
    
    def plot_array(
        self,
        a: np.ndarray,
        masked_values: list = None,
        **kwargs
    ) -> object:
        """Plot 2D array data"""
        ...
    
    def contour_array(
        self,
        a: np.ndarray,
        masked_values: list = None,
        tri_mask: bool = False,
        **kwargs
    ) -> object:
        """Create contour plot of array data"""
        ...
    
    def plot_ibound(
        self,
        ibound: np.ndarray = None,
        color_noflow: str = 'black',
        color_ch: str = 'blue',
        color_vpt: str = 'red',
        color_active: str = None,
        **kwargs
    ) -> object:
        """Plot ibound array"""
        ...
    
    def plot_inactive(
        self,
        ibound: np.ndarray = None,
        color_noflow: str = 'black',
        **kwargs
    ) -> object:
        """Plot inactive cells"""
        ...
    
    def plot_bc(
        self,
        name: str = None,
        package: object = None,
        kper: int = 0,
        color: str = None,
        plotAll: bool = False,
        boundname: str = None,
        **kwargs
    ) -> object:
        """Plot boundary conditions"""
        ...
    
    def plot_shapefile(
        self,
        shp,
        **kwargs
    ) -> object:
        """Plot shapefile data"""
        ...
    
    def plot_shapes(
        self,
        obj,
        **kwargs
    ) -> object:
        """Plot geometric shapes"""
        ...
    
    def plot_centers(
        self,
        **kwargs
    ) -> object:
        """Plot cell centers"""
        ...
    
    def plot_pathline(
        self,
        pl: list = None,
        travel_time: np.ndarray = None,
        **kwargs
    ) -> object:
        """Plot MODPATH pathlines"""
        ...
    
    def plot_endpoint(
        self,
        ep: np.ndarray = None,
        direction: str = 'ending',
        selection: np.ndarray = None,
        **kwargs
    ) -> object:
        """Plot MODPATH endpoints"""
        ...
    
    def plot_timeseries(
        self,
        ts: np.ndarray = None,
        travel_time: float = None,
        **kwargs
    ) -> object:
        """Plot MODPATH timeseries data"""
        ...
    
    def plot_vector(
        self,
        vx: np.ndarray,
        vy: np.ndarray,
        istep: int = 1,
        jstep: int = 1,
        normalize: bool = False,
        masked_values: list = None,
        **kwargs
    ) -> object:
        """Plot vector field"""
        ...
    
    @property
    def extent(self) -> tuple[float, float, float, float]:
        """Plot extent (xmin, xmax, ymin, ymax)"""
        ...
    
    @property
    def mg(self) -> object:
        """Model grid object"""
        ...

class PlotCrossSection:
    """Cross-section plotting for MODFLOW models"""
    def __init__(
        self,
        model: object = None,
        ax: object = None,
        line: dict = None,
        extent: list[float] = None,
        modelgrid: object = None,
        geographic_coords: bool = False,
        **kwargs
    ): ...
    
    def plot_grid(
        self,
        **kwargs
    ) -> object:
        """Plot cross-section grid"""
        ...
    
    def plot_array(
        self,
        a: np.ndarray,
        masked_values: list = None,
        **kwargs
    ) -> object:
        """Plot array data in cross-section"""
        ...
    
    def contour_array(
        self,
        a: np.ndarray,
        masked_values: list = None,
        **kwargs
    ) -> object:
        """Create contour plot in cross-section"""
        ...
    
    def plot_surface(
        self,
        a: np.ndarray,
        masked_values: list = None,
        **kwargs
    ) -> object:
        """Plot surface elevation"""
        ...
    
    def plot_fill_between(
        self,
        a: np.ndarray,
        colors: list = None,
        masked_values: list = None,
        **kwargs
    ) -> object:
        """Plot filled areas between surfaces"""
        ...
    
    def plot_bc(
        self,
        name: str = None,
        package: object = None,
        kper: int = 0,
        color: str = 'red',
        **kwargs
    ) -> object:
        """Plot boundary conditions in cross-section"""
        ...
    
    def plot_ibound(
        self,
        ibound: np.ndarray = None,
        color_noflow: str = 'black',
        color_ch: str = 'blue',
        **kwargs
    ) -> object:
        """Plot ibound in cross-section"""
        ...
    
    def plot_discharge(
        self,
        frf: np.ndarray,
        fff: np.ndarray,
        flf: np.ndarray = None,
        head: np.ndarray = None,
        normalize: bool = False,
        **kwargs
    ) -> object:
        """Plot discharge vectors in cross-section"""
        ...
    
    def plot_pathline(
        self,
        pl: list = None,
        travel_time: np.ndarray = None,
        **kwargs
    ) -> object:
        """Plot pathlines in cross-section"""
        ...
    
    def plot_endpoint(
        self,
        ep: np.ndarray = None,
        direction: str = 'ending',
        **kwargs
    ) -> object:
        """Plot endpoints in cross-section"""
        ...
    
    @property
    def extent(self) -> tuple[float, float, float, float]:
        """Cross-section extent"""
        ...
    
    @property
    def direction(self) -> str:
        """Cross-section direction"""
        ...

class PlotUtilities:
    """General plotting utilities"""
    
    @staticmethod
    def plot_text(
        ax: object,
        model: object = None,
        text_dict: dict = None,
        **kwargs
    ) -> None:
        """Add text to plot"""
        ...
    
    @staticmethod
    def plot_shapefile(
        shp: str,
        ax: object = None,
        **kwargs
    ) -> object:
        """Plot shapefile data"""
        ...
    
    @staticmethod
    def add_colorbar(
        mappable: object,
        ax: object = None,
        **kwargs
    ) -> object:
        """Add colorbar to plot"""
        ...
    
    @staticmethod
    def set_aspect(
        ax: object,
        aspect: str = 'equal',
        **kwargs
    ) -> None:
        """Set plot aspect ratio"""
        ...

class SwiConcentration:
    """Seawater intrusion concentration plotting"""
    def __init__(
        self,
        model: object = None,
        botm: np.ndarray = None,
        istrat: int = 1,
        nu: np.ndarray = None,
        **kwargs
    ): ...
    
    def plot_concentration(
        self,
        zeta: np.ndarray,
        layer: int = None,
        **kwargs
    ) -> object:
        """Plot concentration distribution"""
        ...
    
    def plot_zeta_surfaces(
        self,
        zeta: np.ndarray,
        **kwargs
    ) -> object:
        """Plot zeta surfaces"""
        ...

def plot_shapefile(
    shp: str,
    ax: object = None,
    layer: str = None,
    **kwargs
) -> object:
    """Plot shapefile data on map"""
    ...

def shapefile_extents(shp: str) -> tuple[float, float, float, float]:
    """Get shapefile spatial extents"""
    ...

def plot_cvfd(
    verts: np.ndarray,
    iverts: list,
    xc: np.ndarray = None,
    yc: np.ndarray = None,
    head: np.ndarray = None,
    **kwargs
) -> object:
    """Plot control volume finite difference grid"""
    ...

# Style constants and color maps
styles = {
    'default': {
        'grid_color': 'black',
        'grid_linewidth': 0.5,
        'inactive_color': 'gray',
        'bc_colors': {
            'WEL': 'red',
            'DRN': 'yellow', 
            'RIV': 'blue',
            'GHB': 'cyan',
            'CHD': 'navy'
        }
    },
    'publication': {
        'grid_color': 'black',
        'grid_linewidth': 0.3,
        'font_size': 12,
        'title_size': 14
    },
    'presentation': {
        'grid_color': 'gray',
        'grid_linewidth': 1.0,
        'font_size': 14,
        'title_size': 18
    }
}

def set_style(style_name: str) -> None:
    """Set plotting style"""
    ...

def get_color_palette(name: str) -> list[str]:
    """Get predefined color palette"""
    ...

import flopy
import flopy.utils as fpu
import matplotlib.pyplot as plt
import numpy as np

# Load or create MODFLOW model
mf = flopy.modflow.Modflow.load('model.nam')

# Create map view plot
fig, ax = plt.subplots(figsize=(12, 10))
mapview = flopy.plot.PlotMapView(model=mf, ax=ax, layer=0)

# Plot model grid
mapview.plot_grid(lw=0.5, color='gray')

# Plot boundary conditions
mapview.plot_bc('WEL', color='red', s=50)  # Wells
mapview.plot_bc('RIV', color='blue', lw=3)  # Rivers
mapview.plot_bc('GHB', color='cyan', lw=2)  # General head boundaries

# Plot head results if available
try:
    hds = fpu.HeadFile('model.hds')
    head = hds.get_data(kstpkper=(0, 0))
    
    # Contour heads
    contours = mapview.contour_array(
        head,
        levels=np.linspace(head.min(), head.max(), 20),
        colors='black',
        linewidths=1
    )
    ax.clabel(contours, inline=True, fontsize=8, fmt='%.1f')
    
    # Fill contours with colors
    filled = mapview.plot_array(
        head,
        cmap='viridis',
        alpha=0.7
    )
    cb = plt.colorbar(filled, ax=ax, shrink=0.8)
    cb.set_label('Head Elevation (m)', fontsize=12)
    
    hds.close()
    
except FileNotFoundError:
    print("Head file not found, skipping head contours")

# Plot inactive cells
if hasattr(mf, 'bas6'):
    mapview.plot_inactive(ibound=mf.bas6.ibound.array, color_noflow='gray')

# Add title and labels
ax.set_title('MODFLOW Model - Layer 1 Heads', fontsize=14, fontweight='bold')
ax.set_xlabel('X Coordinate (m)', fontsize=12)
ax.set_ylabel('Y Coordinate (m)', fontsize=12)

# Set equal aspect ratio
ax.set_aspect('equal')

plt.tight_layout()
plt.show()

import flopy
import flopy.utils as fpu
import matplotlib.pyplot as plt
import numpy as np

# Load model and results
mf = flopy.modflow.Modflow.load('model.nam')

# Read budget file for flow data
cbb = fpu.CellBudgetFile('model.cbb')
frf = cbb.get_data(text='FLOW RIGHT FACE')[0]  # Flow right face
fff = cbb.get_data(text='FLOW FRONT FACE')[0]  # Flow front face
flf = cbb.get_data(text='FLOW LOWER FACE')[0]  # Flow lower face (optional)

# Read head data
hds = fpu.HeadFile('model.hds')
head = hds.get_data(kstpkper=(0, 0))

# Create map view
fig, ax = plt.subplots(figsize=(14, 10))
mapview = flopy.plot.PlotMapView(model=mf, ax=ax, layer=0)

# Plot head contours as background
head_contours = mapview.contour_array(
    head,
    levels=15,
    colors='gray',
    alpha=0.5,
    linewidths=0.8
)

# Plot discharge vectors
discharge_plot = mapview.plot_discharge(
    frf, fff, flf,
    head=head,
    normalize=True,  # Normalize vectors for better visualization
    color='blue',
    scale=50,  # Scale factor for arrow size
    headwidth=3,
    headlength=2,
    width=0.002
)

# Plot model grid (light)
mapview.plot_grid(lw=0.3, color='lightgray', alpha=0.5)

# Plot boundary conditions
mapview.plot_bc('WEL', color='red', marker='o', s=100, label='Wells')
mapview.plot_bc('RIV', color='cyan', lw=4, alpha=0.7, label='Rivers')

# Add legend
ax.legend(loc='upper right')

ax.set_title('Groundwater Flow Vectors', fontsize=16, fontweight='bold')
ax.set_xlabel('X Coordinate (m)', fontsize=12)
ax.set_ylabel('Y Coordinate (m)', fontsize=12)

plt.tight_layout()
plt.show()

# Close files
cbb.close()
hds.close()

import flopy
import flopy.utils as fpu
import matplotlib.pyplot as plt
import numpy as np

# Load model
mf = flopy.modflow.Modflow.load('model.nam')

# Define cross-section line
line = {'line': [(0, 2500), (5000, 2500)]}  # West-east cross-section

# Create cross-section plot
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 12))

# First subplot - Head distribution
xs1 = flopy.plot.PlotCrossSection(model=mf, ax=ax1, line=line)

# Read and plot heads
try:
    hds = fpu.HeadFile('model.hds')
    head = hds.get_data(kstpkper=(0, 0))
    
    # Plot head array
    head_plot = xs1.plot_array(
        head,
        cmap='RdYlBu',
        alpha=0.8
    )
    
    # Add colorbar
    cb1 = plt.colorbar(head_plot, ax=ax1, shrink=0.8)
    cb1.set_label('Head Elevation (m)', fontsize=12)
    
    # Contour heads
    head_contours = xs1.contour_array(
        head,
        levels=10,
        colors='black',
        linewidths=1
    )
    ax1.clabel(head_contours, inline=True, fontsize=8)
    
    hds.close()

except FileNotFoundError:
    print("Head file not found")

# Plot grid and model layers
xs1.plot_grid(lw=1, color='black')

# Plot ibound
if hasattr(mf, 'bas6'):
    xs1.plot_ibound(
        ibound=mf.bas6.ibound.array,
        color_noflow='gray',
        alpha=0.5
    )

ax1.set_title('Cross-Section: Head Distribution', fontsize=14, fontweight='bold')
ax1.set_ylabel('Elevation (m)', fontsize=12)

# Second subplot - Hydraulic conductivity
xs2 = flopy.plot.PlotCrossSection(model=mf, ax=ax2, line=line)

# Plot hydraulic conductivity
if hasattr(mf, 'lpf'):
    k_array = mf.lpf.hk.array
    k_plot = xs2.plot_array(
        k_array,
        cmap='plasma',
        norm=plt.Normalize(vmin=k_array.min(), vmax=k_array.max()),
        alpha=0.8
    )
    
    cb2 = plt.colorbar(k_plot, ax=ax2, shrink=0.8)
    cb2.set_label('Hydraulic Conductivity (m/d)', fontsize=12)

# Plot grid
xs2.plot_grid(lw=1, color='black')

# Plot boundary conditions in cross-section
xs2.plot_bc('WEL', color='red', s=100)

ax2.set_title('Cross-Section: Hydraulic Conductivity', fontsize=14, fontweight='bold')
ax2.set_xlabel('Distance along section (m)', fontsize=12)
ax2.set_ylabel('Elevation (m)', fontsize=12)

plt.tight_layout()
plt.show()

import flopy
import flopy.utils as fpu
import matplotlib.pyplot as plt
import numpy as np

# Load model
mf = flopy.modflow.Modflow.load('model.nam')

# Create map view
fig, ax = plt.subplots(figsize=(14, 10))
mapview = flopy.plot.PlotMapView(model=mf, ax=ax, layer=0)

# Plot base head contours
try:
    hds = fpu.HeadFile('model.hds')
    head = hds.get_data(kstpkper=(0, 0))
    
    contours = mapview.contour_array(
        head,
        levels=15,
        colors='lightgray',
        linewidths=1,
        alpha=0.7
    )
    hds.close()
except:
    pass

# Plot model grid (light)
mapview.plot_grid(lw=0.3, color='lightgray', alpha=0.5)

# Plot MODPATH results
try:
    # Read pathline data
    pth_file = fpu.PathlineFile('model.mppth')
    pathlines = pth_file.get_alldata()
    
    # Plot pathlines colored by travel time
    pathline_plot = mapview.plot_pathline(
        pathlines,
        travel_time=None,  # Will use time from pathline data
        lw=2,
        alpha=0.8,
        cmap='viridis'
    )
    
    print(f"Plotted {len(pathlines)} pathlines")
    
    # Read endpoint data
    ep_file = fpu.EndpointFile('model.mpend')
    endpoints = ep_file.get_alldata()
    
    # Plot starting points (green) and ending points (red)
    starting_points = mapview.plot_endpoint(
        endpoints,
        direction='starting',
        marker='o',
        color='green',
        s=50,
        label='Starting points',
        alpha=0.8
    )
    
    ending_points = mapview.plot_endpoint(
        endpoints,
        direction='ending',
        marker='s',
        color='red',
        s=30,
        label='Ending points',
        alpha=0.8
    )
    
    print(f"Plotted {len(endpoints)} particle endpoints")
    
except FileNotFoundError as e:
    print(f"MODPATH files not found: {e}")

# Plot boundary conditions
mapview.plot_bc('WEL', color='blue', marker='o', s=150, 
               edgecolor='darkblue', linewidth=2, label='Wells')
mapview.plot_bc('RIV', color='cyan', lw=5, alpha=0.7, label='Rivers')

# Add colorbar for pathlines if they were plotted
try:
    if 'pathline_plot' in locals():
        cb = plt.colorbar(pathline_plot, ax=ax, shrink=0.8)
        cb.set_label('Travel Time (days)', fontsize=12)
except:
    pass

# Legend
ax.legend(loc='upper right', fontsize=10)

ax.set_title('Particle Tracking Results', fontsize=16, fontweight='bold')
ax.set_xlabel('X Coordinate (m)', fontsize=12)
ax.set_ylabel('Y Coordinate (m)', fontsize=12)
ax.set_aspect('equal')

plt.tight_layout()
plt.show()

import flopy
import flopy.utils as fpu
import matplotlib.pyplot as plt
import numpy as np

# Load model
mf = flopy.modflow.Modflow.load('model.nam')
nlay = mf.dis.nlay

# Create multi-panel figure
fig, axes = plt.subplots(2, 2, figsize=(16, 12))
axes = axes.flatten()

# Read head data
try:
    hds = fpu.HeadFile('model.hds')
    head = hds.get_data(kstpkper=(0, 0))
    hds.close()
    
    # Global head range for consistent coloring
    vmin, vmax = head.min(), head.max()
    levels = np.linspace(vmin, vmax, 20)
    
except FileNotFoundError:
    head = None
    vmin = vmax = None
    levels = None

# Plot each layer
for i in range(min(nlay, 4)):  # Up to 4 layers
    ax = axes[i]
    mapview = flopy.plot.PlotMapView(model=mf, ax=ax, layer=i)
    
    # Plot grid
    mapview.plot_grid(lw=0.3, color='black', alpha=0.5)
    
    # Plot head data if available
    if head is not None:
        head_plot = mapview.plot_array(
            head,
            cmap='RdYlBu',
            vmin=vmin,
            vmax=vmax,
            alpha=0.8
        )
        
        # Contour lines
        contours = mapview.contour_array(
            head,
            levels=levels[::2],  # Every other contour
            colors='black',
            linewidths=0.8,
            alpha=0.7
        )
    
    # Plot boundary conditions
    mapview.plot_bc('WEL', color='red', s=30, alpha=0.8)
    mapview.plot_bc('RIV', color='blue', lw=2, alpha=0.7)
    
    # Plot inactive cells
    if hasattr(mf, 'bas6'):
        mapview.plot_inactive(
            ibound=mf.bas6.ibound.array,
            color_noflow='lightgray'
        )
    
    ax.set_title(f'Layer {i+1}', fontsize=12, fontweight='bold')
    ax.set_xlabel('X (m)', fontsize=10)
    ax.set_ylabel('Y (m)', fontsize=10)
    ax.set_aspect('equal')

# Add shared colorbar
if head is not None and 'head_plot' in locals():
    # Create colorbar axis
    cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
    cb = plt.colorbar(head_plot, cax=cbar_ax)
    cb.set_label('Head Elevation (m)', fontsize=12, rotation=270, labelpad=20)

fig.suptitle('Multi-Layer Head Distribution', fontsize=16, fontweight='bold')
plt.tight_layout()
plt.subplots_adjust(right=0.9)
plt.show()

import flopy
import flopy.utils as fpu
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LogNorm

# Load MT3DMS results
try:
    ucn = fpu.UcnFile('MT3D001.UCN')
    times = ucn.get_times()
    
    # Create time series plot
    fig = plt.figure(figsize=(16, 12))
    
    # Select representative times
    time_indices = [0, len(times)//4, len(times)//2, -1]
    
    for idx, time_idx in enumerate(time_indices):
        ax = plt.subplot(2, 2, idx + 1)
        
        # Get concentration data
        conc = ucn.get_data(idx=time_idx)
        time_val = times[time_idx]
        
        # Create map view (assuming structured grid)
        # Note: Would need model object for proper grid, using generic approach
        
        # Plot concentration with log scale for better visualization
        conc_masked = np.ma.masked_where(conc <= 0, conc)
        
        im = ax.imshow(
            conc_masked[0, :, :],  # Top layer
            cmap='plasma',
            norm=LogNorm(vmin=0.01, vmax=conc.max()),
            origin='lower',
            alpha=0.8
        )
        
        # Add contour lines
        contour_levels = np.logspace(-2, np.log10(conc.max()), 8)
        cs = ax.contour(
            conc[0, :, :],
            levels=contour_levels,
            colors='black',
            linewidths=1,
            alpha=0.6
        )
        ax.clabel(cs, inline=True, fontsize=8, fmt='%.2f')
        
        ax.set_title(f'Time = {time_val:.1f} days', fontsize=12, fontweight='bold')
        ax.set_xlabel('Column', fontsize=10)
        ax.set_ylabel('Row', fontsize=10)
        ax.set_aspect('equal')
    
    # Add shared colorbar
    cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
    cb = plt.colorbar(im, cax=cbar_ax)
    cb.set_label('Concentration (mg/L)', fontsize=12, rotation=270, labelpad=20)
    
    fig.suptitle('Contaminant Plume Evolution', fontsize=16, fontweight='bold')
    plt.tight_layout()
    plt.subplots_adjust(right=0.9)
    plt.show()
    
    ucn.close()

except FileNotFoundError:
    print("MT3DMS concentration file not found")

# Concentration breakthrough curves
try:
    ucn = fpu.UcnFile('MT3D001.UCN')
    
    # Define monitoring points
    monitor_points = [(0, 10, 25), (0, 20, 50), (0, 30, 75)]
    
    # Get time series
    conc_ts = ucn.get_ts(idx=monitor_points)
    times = ucn.get_times()
    
    # Plot breakthrough curves
    plt.figure(figsize=(12, 8))
    
    for i, (lay, row, col) in enumerate(monitor_points):
        plt.semilogy(
            times,
            conc_ts[:, i],
            'o-',
            linewidth=2,
            markersize=4,
            label=f'MW-{i+1} (R{row}, C{col})'
        )
    
    plt.xlabel('Time (days)', fontsize=12)
    plt.ylabel('Concentration (mg/L)', fontsize=12)
    plt.title('Concentration Breakthrough Curves', fontsize=14, fontweight='bold')
    plt.legend(fontsize=10)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()
    
    ucn.close()

except FileNotFoundError:
    print("MT3DMS concentration file not found")

import flopy
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np

# Set publication style
plt.style.use('seaborn-v0_8-paper')  # Clean, publication-ready style
mpl.rcParams['font.size'] = 11
mpl.rcParams['axes.labelsize'] = 12
mpl.rcParams['axes.titlesize'] = 14
mpl.rcParams['xtick.labelsize'] = 10
mpl.rcParams['ytick.labelsize'] = 10
mpl.rcParams['legend.fontsize'] = 10
mpl.rcParams['figure.titlesize'] = 16

# Custom color schemes
custom_colors = {
    'wells': '#d62728',      # Red
    'rivers': '#1f77b4',     # Blue
    'ghb': '#ff7f0e',        # Orange
    'drn': '#2ca02c',        # Green
    'inactive': '#7f7f7f',   # Gray
    'grid': '#cccccc'        # Light gray
}

# Load model
mf = flopy.modflow.Modflow.load('model.nam')

# Create publication figure
fig, ax = plt.subplots(figsize=(10, 8), dpi=300)  # High DPI for publication

# Create map view
mapview = flopy.plot.PlotMapView(model=mf, ax=ax, layer=0)

# Plot with custom styling
mapview.plot_grid(
    lw=0.5,
    color=custom_colors['grid'],
    alpha=0.7
)

# Plot boundary conditions with custom styling
mapview.plot_bc(
    'WEL',
    color=custom_colors['wells'],
    marker='o',
    s=80,
    edgecolor='darkred',
    linewidth=1.5,
    label='Pumping Wells',
    zorder=5
)

mapview.plot_bc(
    'RIV',
    color=custom_colors['rivers'],
    lw=3,
    alpha=0.8,
    label='Rivers',
    zorder=4
)

mapview.plot_bc(
    'GHB',
    color=custom_colors['ghb'],
    lw=2,
    alpha=0.8,
    label='General Head Boundary',
    zorder=3
)

# Plot head data with custom colormap
try:
    hds = fpu.HeadFile('model.hds')
    head = hds.get_data(kstpkper=(0, 0))
    
    # Use a perceptually uniform colormap
    head_plot = mapview.plot_array(
        head,
        cmap='viridis',
        alpha=0.7,
        zorder=1
    )
    
    # Clean contours
    contours = mapview.contour_array(
        head,
        levels=np.linspace(head.min(), head.max(), 15),
        colors='white',
        linewidths=0.8,
        alpha=0.9,
        zorder=2
    )
    
    # Format colorbar
    cb = plt.colorbar(
        head_plot,
        ax=ax,
        shrink=0.8,
        aspect=20,
        pad=0.02
    )
    cb.set_label('Hydraulic Head (m above MSL)', fontsize=12, labelpad=15)
    cb.ax.tick_params(labelsize=10)
    
    hds.close()
    
except FileNotFoundError:
    pass

# Plot inactive cells
if hasattr(mf, 'bas6'):
    mapview.plot_inactive(
        ibound=mf.bas6.ibound.array,
        color_noflow=custom_colors['inactive'],
        alpha=0.3
    )

# Professional formatting
ax.set_xlabel('Easting (m)', fontsize=12, fontweight='bold')
ax.set_ylabel('Northing (m)', fontsize=12, fontweight='bold')
ax.set_title('Groundwater Flow Model - Steady State Heads', 
             fontsize=14, fontweight='bold', pad=20)

# Clean legend
legend = ax.legend(
    loc='upper left',
    frameon=True,
    fancybox=True,
    shadow=True,
    fontsize=10
)
legend.get_frame().set_facecolor('white')
legend.get_frame().set_alpha(0.9)

# Set aspect ratio and tight layout
ax.set_aspect('equal')

# Add north arrow and scale bar (custom functions would be needed)
# add_north_arrow(ax, xy=(0.9, 0.9), size=0.05)
# add_scale_bar(ax, length=1000, location='lower right')

# Grid and ticks
ax.grid(True, alpha=0.3, linestyle='--', linewidth=0.5)
ax.tick_params(axis='both', which='major', labelsize=10)

# Tight layout with padding
plt.tight_layout(pad=1.5)

# Save in multiple formats for publication
plt.savefig('model_heads.png', dpi=300, bbox_inches='tight', 
            facecolor='white', edgecolor='none')
plt.savefig('model_heads.pdf', bbox_inches='tight', 
            facecolor='white', edgecolor='none')
plt.savefig('model_heads.svg', bbox_inches='tight', 
            facecolor='white', edgecolor='none')

plt.show()

# Plotting types
PlotArray = np.ndarray  # 2D or 3D array for plotting
ColorMap = Union[str, mpl.colors.Colormap]
PlotExtent = tuple[float, float, float, float]  # (xmin, xmax, ymin, ymax)

# Line and geometry types
CrossSectionLine = dict[str, list[tuple[float, float]]]
PlotLine = dict[str, Union[list, np.ndarray]]
Coordinates = tuple[float, float]

# Style and formatting types
PlotStyle = dict[str, Union[str, int, float]]
ColorSpec = Union[str, tuple[float, float, float]]
MarkerStyle = str
LineStyle = str

# Vector data types
VectorField = tuple[np.ndarray, np.ndarray]  # (vx, vy)
FlowData = tuple[np.ndarray, np.ndarray, np.ndarray]  # (frf, fff, flf)
DischargeData = tuple[np.ndarray, np.ndarray, np.ndarray]  # (qx, qy, qz)

# Particle tracking types
PathlineData = list[np.ndarray]
EndpointData = np.ndarray
TimeseriesData = np.ndarray

# 3D plotting types
SurfaceData = np.ndarray  # 2D array for surface plots
VoxelData = np.ndarray   # 3D boolean array for voxel plots
ScatterData = tuple[np.ndarray, np.ndarray, np.ndarray]  # (x, y, z)

# File and output types
ImageFormat = Literal['png', 'pdf', 'svg', 'eps', 'jpg', 'tiff']
PlotSize = tuple[float, float]  # (width, height) in inches
DPI = int  # Dots per inch for raster output