%matplotlib inline
import numpy as np

"""
Explanation: <a href="http://landlab.github.io"><img style="float: left" src="../../landlab_header.png"></a>
Using the Landlab flexure component
<hr>
<small>For more Landlab tutorials, click here: <a href="https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html">https://landlab.readthedocs.io/en/latest/user_guide/tutorials.html</a></small>
<hr>

In this example we will:
* create a Landlab component that solves the two-dimensional elastic flexure equation
* apply randomly distributed point loads
* run the component
* plot some output
A bit of magic so that we can plot within this notebook.
End of explanation
"""


from landlab import RasterModelGrid
from landlab.plot.imshow import imshow_grid

"""
Explanation: Create the grid
We are going to build a uniform rectilinear grid with a node spacing of 10 km in the y-direction and 20 km in the x-direction on which we will solve the flexure equation.
First we need to import RasterModelGrid.  We also import the Landlab plotting function imshow_grid to view the grids.
End of explanation
"""


grid = RasterModelGrid((200, 400), xy_spacing=(10e3, 20e3))

grid.dy, grid.dx

"""
Explanation: Create a rectilinear grid with a spacing of 10 km between rows and 20 km between columns. The numbers of rows and columms are provided as a tuple of (n_rows, n_cols), in the same manner as similar numpy functions. The spacing is also a tuple, (dy, dx).
End of explanation
"""


from landlab.components.flexure import Flexure

"""
Explanation: Create the component
Now we create the flexure component and tell it to use our newly-created grid. First, though, we'll examine the Flexure component a bit.
End of explanation
"""


Flexure.input_var_names

"""
Explanation: The Flexure component, as with most landlab components, will require our grid to have some data that it will use. We can get the names of these data fields with the intput_var_names attribute of the component class.
End of explanation
"""


Flexure.var_units('lithosphere__overlying_pressure_increment')

"""
Explanation: We see that flexure uses just one data field: the change in lithospheric loading. landlab component classes can provide additional information about each of these fields. For instance, to the the units for a field, use the var_units method.
End of explanation
"""


Flexure.var_help('lithosphere__overlying_pressure_increment')

"""
Explanation: To print a more detailed description of a field, use var_help.
End of explanation
"""


Flexure.output_var_names

Flexure.var_help('lithosphere_surface__elevation_increment')

"""
Explanation: What about the data that Flexure provides? Use the output_var_names attribute.
End of explanation
"""


grid.add_zeros("lithosphere__overlying_pressure_increment", at="node")
flex = Flexure(grid, method='flexure', n_procs=4)

"""
Explanation: Now that we understand the component a little more, create it using our grid.
End of explanation
"""


load = np.random.normal(0, 100 * 2650. * 9.81, grid.number_of_nodes)
grid.at_node['lithosphere__overlying_pressure_increment'] = load

imshow_grid(grid,
            'lithosphere__overlying_pressure_increment',
            symmetric_cbar=True,
            cmap='nipy_spectral')

"""
Explanation: Add some loading
We will add loads to the grid. As we saw above, for this component, the name of the variable that holds the applied loads is lithosphere__overlying_pressure. We add loads of random magnitude at every node of the grid.
End of explanation
"""


flex.update()

"""
Explanation: Update the component to solve for deflection
If you have more than one processor on your machine you may want to use several of them.
End of explanation
"""


imshow_grid(grid,
            'lithosphere_surface__elevation_increment',
            symmetric_cbar=True,
            cmap='nipy_spectral')

"""
Explanation: As we saw above, the flexure component creates an output field (lithosphere_surface__elevation_increment) that contains surface deflections for the applied loads.
Plot the output
We now plot these deflections with the imshow_grid method, which is available to all landlab components.
End of explanation
"""


flex.eet *= 2.
flex.update()
imshow_grid(grid,
            'lithosphere_surface__elevation_increment',
            symmetric_cbar=True,
            cmap='nipy_spectral')

"""
Explanation: Maintain the same loading distribution but double the effective elastic thickness.
End of explanation
"""


load[np.where(np.logical_and(grid.node_x>3000000, grid.node_x<5000000))]= \
    load[np.where(np.logical_and(grid.node_x>3000000, grid.node_x<5000000))]+1e7
imshow_grid(grid,
            'lithosphere__overlying_pressure_increment',
            symmetric_cbar=True,
            cmap='nipy_spectral')

flex.update()
imshow_grid(grid,
            'lithosphere_surface__elevation_increment',
            symmetric_cbar=True,
            cmap='nipy_spectral')

"""
Explanation: Now let's add a vertical rectangular load to the middle of the grid.  We plot the load grid first to make sure we did this correctly.
End of explanation
"""


