# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Gridtools Module
#
# A new gridtools.py module has some tools for reading in a solution as a set of grid patches computed using AMR and extracting data on a uniform 2d grid or along a 1d transect, by interpolating from the finest level patch available at each point.
#
# You can read in a time frame of the solution at some fixed time using, e.g.:
# ```
# from clawpack.pyclaw import solution
# framesoln = solution.Solution(frameno, path=outdir, file_format='binary')
# ```
# and then pass `framesoln` in to `gridtools.grid_output_2d` along with other arguments that specify what set of scalar values to extract and the set of grid points on which to extract them:
# ```
# def grid_output_2d(framesoln, out_var, xout, yout, levels='all', 
#                    return_ma=True):
#     :Input:
#         framesoln:  One frame of Clawpack solution (perhaps with AMR),
#                  An object of type pyclaw.Solution.solution.
#         out_var: function that maps q to desired quantities Q[m,i,j] or
#                  Q[i,j] if only one.
#                  If type(out_var) == int, then Q[i,j] = q[out_var,i,j]
#         xout, yout: arrays of output points (1d or 2d arrays)
#         levels: list of levels to use, or 'all'
#         return_ma: True to return as masked_array, False to return with
#                 NaN in locations that framesoln doesn't cover.
#     :Output:
#         qout: Solution obtained on xout,yout grid
# ```

# %matplotlib inline

from pylab import *
from IPython.display import Image
import os,sys
from clawpack.visclaw import colormaps, frametools, geoplot
from clawpack.geoclaw import dtopotools
from clawpack.pyclaw.solution import Solution
from clawpack.pyclaw import solution

sys.path.insert(0,'../new_python')
import region_tools, topotools, marching_front
from plottools import pcolorcells
import gridtools

# ## Sample data from test case
#
# We use output that can be generated by running the notebook [RunGeoclaw_test1.ipynb](RunGeoclaw_test1.ipynb) in the directory `examples/geoclaw_test1`.  See that notebook for more discussion of this test problem.

rundir = '../examples/geoclaw_test1'
outdir = os.path.join(rundir, '_output_3')
sys.path.insert(0,rundir)

frameno = 5
framesoln = solution.Solution(frameno, path=outdir, file_format='binary')

print('Frame %i solution at t = %.0f seconds has %i grid patches' \
      % (frameno, framesoln.t, len(framesoln.states)))

# ### Plot using frametools
#
# The standard way to plot the AMR solution using visclaw is to provide a `setplot.py` file that specifies the desired plots and then use `clawpack.visclaw.frametools` to loop over all the grid patches and produce the desired plots.  This is invoked behind the scenes when doing `make plots` or using the interactive `Iplotclaw` module. But it is also possible to use frametools directly to produce one set of plots, for example:

from setplot import setplot

plotdata = setplot()
plotdata.outdir = outdir
frametools.plotframe(frameno,plotdata)

# ## Extract uniform grid
#
# But now suppose we want to extract values on a uniform 2D grid for some purpose, e.g. when making an animation over some region.  
#
# The water surface eta is given by `q[3,:,:]` and the topography B can be computed by subtracting the water depth `q[0,:,:]` from this, so we can define two functions that return these as 2D arrays for any `q` defining the full solution on a grid patch:

eta = lambda q: q[3,:,:]
B = lambda q: q[3,:,:]-q[0,:,:]

# Define the desired output grid.  For illustration we use a very coarse grid:

x = linspace(-0.005,0.01,16)
y = linspace(-0.01,0.01,21)
Xout, Yout = meshgrid(x,y)

B_out_2d = gridtools.grid_output_2d(framesoln, B, Xout, Yout, 
                                    levels='all',return_ma=True)
eta_out_2d = gridtools.grid_output_2d(framesoln, eta, Xout, Yout, 
                                      levels='all',return_ma=True)
print('Interpolated to uniform grids of shape ', B_out_2d.shape)

# +
figure(figsize=(10,6))
pcolorcells(Xout, Yout, B_out_2d, cmap=geoplot.land_colors)
clim(-0.5,0.5)
contour(Xout, Yout, B_out_2d, [0], colors='k')

h_out_2d = eta_out_2d - B_out_2d
eta_masked = ma.masked_where(h_out_2d < 0.001, eta_out_2d)
pcolorcells(Xout, Yout, eta_masked, cmap=geoplot.tsunami_colormap)
clim(-0.5,0.5)
colorbar()
gca().set_aspect(1)
title('Surface on uniform coarse grid\nBlack contour is B=0 at this resolution');
# -

# ## Extract 1d transects
#
# It is often difficult to visualize the topography and water depth from 2d plots like those shown above, and so it is useful to plot the solution along 1d transects. 
#
# As an example, we plot the solution along a transect at constant latitude `y = 0.002` over `-0.005 <= x <= 0.01`, which goes through the Gaussian depression near the shore.
#
# We also illustrate that a single call to `gridtools.grid_output_2d` can be used for each frame by defining `out_var` below to be an array that will return both `B_out` and `eta_out`.  This is more efficient for large data sets and several output quantities than multiple calls to `gridtools.grid_output_2d`.

eta = lambda q: q[3,:,:]
B = lambda q: q[3,:,:]-q[0,:,:]
out_var = lambda q: array((B(q),eta(q)))

# +
# output grid (1d transect):
#xout = linspace(-0.005, 0.01, 1001)
xout = linspace(-0.005, 0.03, 1001)
ylat = 0.002
yout = ylat * ones(xout.shape)

# single call to extract both quantities of interest:
qout = gridtools.grid_output_2d(framesoln, out_var, xout, yout, 
                                 levels='all',return_ma=True)

# unpack the results:
B_out = qout[0,:]
eta_out = qout[1,:]
# -

# Plot the transect results, using `fill_between` to show the cross section of earth as green and of water as blue:

figure(figsize=(10,4))
fill_between(xout, eta_out, B_out, color=[.5,.5,1])
fill_between(xout, B_out, -6, color=[.7,1,.7])
plot(xout, B_out, 'g')
plot(xout, eta_out, 'b')

# Note that we are interpolating to a fine grid with 1001 points, and piecewise constant interpolation is performed using the cell average values. So in the plot above the curves look piecewise constant with jumps at the cell interfaces of the computational grid from which the solution is interpolated.

# ### Loop over frames
#
# Putting this in a loop lets us see much better how the solution evolves along the coast.
#
# For these plots we zoom in on the region near the coast.
#
# Note in the plots below that at early times only a coarse grid is present in this region, and the interpolated solution clearly shows this coarse grid structure.
#
# Also note that we are plotting results from the version of this example in which the `force_dry` mask is used to indicate cells that should be initialized to dry (`h = 0`) even if the topography is below sea level (`B < 0`). However, this is applied only on the finest grid and so at early times there is water in the depression that disappears at time 600, when the finest grid is introduced (which has been carefully chosen to be before the tsunami arrives). 

# +
xout = linspace(-0.005, 0.01, 1001)
ylat = 0.002
yout = ylat * ones(xout.shape)

for frameno in range(6):
    framesoln = solution.Solution(frameno, path=outdir, file_format='binary')
    qout = gridtools.grid_output_2d(framesoln, out_var, xout, yout, 
                                     levels='all',return_ma=True)
    B_out = qout[0,:]
    eta_out = qout[1,:]
    figure(figsize=(10,4))
    fill_between(xout, eta_out, B_out, color=[.5,.5,1])
    fill_between(xout, B_out, -6, color=[.7,1,.7])
    plot(xout, B_out, 'g')
    plot(xout, eta_out, 'b')
    title('Transect along y = %.4f at t = %.1f' % (ylat, framesoln.t))
# -

# ## Transects at an angle to the grid
#
# In the example above our transect was along a line of constant latitude, but this is not necessary.  A transect between any two points `(x1,y1)` and `(x2,y2)` can be defined by e.g.

x1 = -0.004; x2 = 0.008
y1 = -0.005; y2 = 0.0075
npts = 1001
xout = linspace(x1,x2,npts)
yout = linspace(y1,y2,npts)

# +
figure(figsize=(10,6))
pcolorcells(Xout, Yout, B_out_2d, cmap=geoplot.land_colors)
clim(-0.5,0.5)
contour(Xout, Yout, B_out_2d, [0], colors='k')

h_out_2d = eta_out_2d - B_out_2d
eta_masked = ma.masked_where(h_out_2d < 0.001, eta_out_2d)
pcolorcells(Xout, Yout, eta_masked, cmap=geoplot.tsunami_colormap)
clim(-0.5,0.5)
colorbar()
gca().set_aspect(1)

plot(xout,yout,'w',linewidth=2)
text(0.006,0.008,'Transect',color='w',fontsize=15)
title('Surface on uniform coarse grid\nBlack contour is B=0 at this resolution');

# +
qout = gridtools.grid_output_2d(framesoln, out_var, xout, yout, 
                                 levels='all',return_ma=True)
B_out = qout[0,:]
eta_out = qout[1,:]

figure(figsize=(10,4))
fill_between(xout, eta_out, B_out, color=[.5,.5,1])
fill_between(xout, B_out, -6, color=[.7,1,.7])
plot(xout, B_out, 'g')
plot(xout, eta_out, 'b')
xlabel('Longitude x')
title('Plot cross section on transect vs longitude');
# -

# (Note that the 2d plot above showed the coarser resolution uniform grid solution extracted above, while the transect plot uses the full AMR solution.)

# ### Plot vs. distance along transect
#
# In the plot above we plotted the value on the transect vs. longitude.  If the transect had been running more N-S than E-W we could have instead plotted against latitude.
#
# Sometimes we want to plot values on the transect vs. the distance in meters.  When GeoClaw is used in longitude-latitude coordinates, this distance can be calculated using the `clawpack.geoclaw.util.haversine` function:

from clawpack.geoclaw import util
dist = util.haversine(x1, y1, xout, yout)
print('The length of this transect is %.2f meters' % dist[-1])

figure(figsize=(10,4))
fill_between(dist, eta_out, B_out, color=[.5,.5,1])
fill_between(dist, B_out, -6, color=[.7,1,.7])
plot(dist, B_out, 'g')
plot(dist, eta_out, 'b')
xlabel('Distance along transect (meters)')
title('Plot cross section on transect vs distance');



StarcoderdataJupyterScriptsDedupFiltered