#!/usr/bin/env python
# coding: utf-8

# # Evolution of a fan
# 
# This notebook reproduces the [fan example](https://fastscape-lem.github.io/fastscapelib-fortran/#_fan_f90) provided in the fastscapelib-fortran library. It illustrates continental transport/deposition.

# In[ ]:


import numpy as np
import xsimlab as xs
import matplotlib.pyplot as plt
import fastscape

get_ipython().run_line_magic('matplotlib', 'inline')


# In[ ]:


print('xarray-simlab version: ', xs.__version__)
print('fastscape version: ', fastscape.__version__)


# ## Import, customize and inspect the model
# 
# A sediment model is available in [fastscape](https://fastscape.readthedocs.io/en/latest/).

# In[ ]:


from fastscape.models import sediment_model


# We want to model the evolution of a passive escarpment, thus with no uplift.

# In[ ]:


from fastscape.processes import Escarpment


model = (sediment_model
         .drop_processes('uplift')
         .update_processes({'init_topography': Escarpment}))


# In[ ]:


model


# In[ ]:


model.visualize(show_inputs=True)


# ## Model setup

# In[ ]:


in_ds = xs.create_setup(
    model=model,
    clocks={
        'time': np.arange(0, 4e5 + 2e3, 2e3),
        'out': np.arange(0, 4e5 + 4e3, 4e3),
    },
    master_clock='time',
    input_vars={
        'grid__shape': [101, 201],
        'grid__length': [1e4, 2e4],
        'boundary__status': ['fixed_value', 'core', 'looped', 'looped'],
        'init_topography': {
            'x_left': 1e4,
            'x_right': 1e4,
            'elevation_left': 0.,
            'elevation_right': 1e3
        },
        'flow__slope_exp': 1.,
        'spl': {
            'k_coef_bedrock': 1e-4,
            'k_coef_soil': 1.5e-4,
            'g_coef_bedrock': 1.,
            'g_coef_soil': 1.,
            'area_exp': 0.4,
            'slope_exp': 1.
        },
        'diffusion': {
            'diffusivity_bedrock': 1e-2,
            'diffusivity_soil': 1.5e-2
        }
    },
    output_vars={
        'topography__elevation': 'out',
        'erosion__rate': 'out'
    }
)

in_ds


# ## Run the model
# 

# In[ ]:


with xs.monitoring.ProgressBar():
    out_ds = in_ds.xsimlab.run(model=model)


# In[ ]:


out_ds


# ## Plot the outputs
# 
# The following plots (total erosion/deposition rate, drainage area and topography cross-section) highlight the dynamics on the left area (depositional) and its interaction with the right area (erosional).

# In[ ]:


import hvplot.xarray
import holoviews as hv
from xshade import hillshade


erosion_plot = out_ds.erosion__rate.hvplot.image(
    x='x', y='y', clim=(-2.5e-3, 2.5e-3),
    cmap='RdYlGn_r', groupby='out'
)

hillshade_plot = hillshade(out_ds, 'out').hvplot.image(
    x='x', y='y', cmap='gray', alpha=0.4,
    colorbar=False, hover=False, groupby='out'
)

ysections =  out_ds.topography__elevation.sel(y=[2.5e3, 5e3, 7.5e3])

sections_plot = ysections.hvplot.line(
    by='y', groupby='out', ylim=(0, 1100),
    height=200, legend='top_left',
)

hv.Layout((erosion_plot * hillshade_plot) + sections_plot).cols(1)


# In[ ]: