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

# ## Visualizing IRIS netcdf tomography models
# 
# This notebook demonstrates two ways of plotting seismic data from IRIS-EMC tomographic models:
# 
# 1. [Volume Rendering](#Volume-Rendering)
# 2. [Fixed Depth Maps](#Fixed-Depth-Maps)
# 
# 
# ## Volume Rendering 
# 
# 
# The volume rendering here is based on the notebook presented at [Havlin et al. 2020](https://github.com/earthcube2020/ec20_havlin_etal). In order to use the volume renderer, we'll need to interpolate the IRIS netcdf models from their geographic coordinates to cartesian coordinates, for which we'll use the https://github.com/chrishavlin/yt_velmodel_vis repository. The notebook from Havlin et al. 2020 describes the details of the interpolations and some additional annotations for volume rendering contained in the `yt_velmodel_vis` repository, so for the present notebook, we've moved many of those function calls to the `seismic_helper` module in the `resources` directory for the present notebook. See the EarthCube notebook for a detailed explanation.
# 
# First, we load an IRIS EMC model into memory, use the `yt_velmodel_vis.seis_model` class to interpolate to cartesian and then load a *yt* dataset using `yt.load_uniform_grid`. We'll use shear wave velocity anomalies in the Western U.S. from James et al., 2011 (avaialable from IRIS [here](https://http://ds.iris.edu/ds/products/emc-nwus11-s/)). 

# In[1]:


# imports and initialization
import os
import yt
import matplotlib.pyplot as plt
from yt_velmodel_vis import seis_model as SM, transferfunctions as TFs
from resources import seismic_helper as SH 

# load interpolated data using the yt uniform grid loader (not editable)
# set the model file and the field to visualize
modelfile='NWUS11-S_percent.nc' # the model netCDF file 
datafld='dvs' # the field to visualize, must be a variable in the netCDF file

# set the interpolation dictionary. If the interpolation for this model does 
# not exist, SM.netcdf() will build it. 
interp_dict={'field':datafld,'max_dist':50000,'res':[10000,10000,10000],
              'input_units':'m','interpChunk':int(1e7)}

# load the model 
model=SM.netcdf(modelfile,interp_dict)

# set some objects required for loading in yt 
bbox = model.cart['bbox'] # the bounding box of interpolated cartesian grid
data={datafld:model.interp['data'][datafld]} # data container for yt scene

# load the data as a uniform grid, create the 3d scene
ds = yt.load_uniform_grid(data,data[datafld].shape,1.0,bbox=bbox,nprocs=1,
                        periodicity=(True,True,True),unit_system="mks")


# Volume rendering in *yt* requires that we specify a `Scene` comprised of `sources` representing the volumes of data that we wish to render, a `transfer_function` that controls how a volume is sampled, and `camera` settings that control orientation. The documentation on volume rendering [here](https://yt-project.org/doc/visualizing/volume_rendering.html) describes these components in more detail. 
# 
# First, we set up a transfer function for the volume rendering following [Havlin et al. 2020](https://github.com/earthcube2020/ec20_havlin_etal). This transfer function highlights a region of slow and fast anomalies, linearly varying the opacity and color within the slow and fast segments to emphasize interior structure: 

# In[2]:


# setting up transfer functions 
tfOb = TFs.dv(data[datafld].ravel(),bounds=[-4,4])

# segment 1, slow anomalies
bnds=[-1.3,-.3]
TFseg=TFs.TFsegment(tfOb,bounds=bnds,cmap='OrRd_r')
alpha_o=0.95
Dalpha=-0.85
alpha=alpha_o + Dalpha/(bnds[1]-bnds[0]) * (TFseg.dvbins_c-bnds[0])
tfOb.addTFsegment(alpha,TFseg)

# segment 2, fast anomalies
bnds=[.1,.35]
TFseg=TFs.TFsegment(tfOb,bounds=bnds,cmap='winter_r')
alpha_o=.6
Dalpha=.4
alpha=alpha_o+ Dalpha/(bnds[1]-bnds[0]) * (TFseg.dvbins_c-bnds[0])
tfOb.addTFsegment(alpha,TFseg)
    
SH.plotTf(tfOb)
print("Ready to build scene")


# Now we build our scene using the `configure_scene` helper function, which sets the transfer function, camera orientation and adds annotations.

# In[6]:


sc = SH.configure_scene(ds, datafld, model, bbox, tfOb.tf,res_factor=1.25)


# The `configure_scene` helper returns a scene that will render when we call `sc.show()`: 

# In[7]:


sc.show(sigma_clip=1.5)


# The annotations added by `configure_scene` include state boundaries, plate boundaries and recent volcanos. 

# In[10]:


sc.save('./figures/seismic_render.png', sigma_clip = 1.5, render=False)


# ### Fixed Depth Maps
# 
# *yt* supports geographic transforms and plotting via Cartopy ([link to *yt* docs](https://yt-project.org/doc/visualizing/geographic_projections_and_transforms.html)). To take advantage of the geographic transforms, we can load the IRIS netcdf fiels directly into *yt* as uniform gridded data with the `internal_geographic` geometry specification: 
# 
# ```python
# ds = yt.load_uniform_grid(data, sizes, 1.0, geometry=("internal_geographic", dims),bbox=bbox)
# ```
# 
# Before doing so, we first need to load our netcdf data into an in-memory flat `data` dictionary, which we do here following the *yt* example [notebook for using geographic projections](https://yt-project.org/doc/cookbook/geographic_projections.html#cookbook-geographic-projections).

# In[1]:


import numpy as np
import yt 
import os
import re
import cartopy.crs as ccrs
import cartopy.feature as cf
import netCDF4 as nc4

def get_data_path(arg):
    if os.path.exists(arg):
        return arg
    else:
        return os.path.join(yt.config.ytcfg.get("yt", "test_data_dir"), arg)


# In[2]:


w_regex = re.compile(r'([a-zA-Z]+)(.*)')
def regex_parser(s):
    try:
        return "**".join(filter(None, w_regex.search(s).groups()))
    except AttributeError:
        return s

def get_dims(n):
    dims = []
    sizes = []
    bbox = []
    ndims = len(n.dimensions)
    for dim in n.dimensions.keys():
        size = n.variables[dim].size
        if size > 1:
            bbox.append([n.variables[dim][:].min(),
                         n.variables[dim][:].max()])
            dims.append(n.variables[dim].name)
            sizes.append(size)
   
    bbox = np.array(bbox)
    return dims, bbox, sizes, ndims
    
def prep_netcdf(n):
    
    dims, bbox, sizes, ndims = get_dims(n)
    
    data = {}
    names = {}
    for field, d in n.variables.items():
        if d.ndim != ndims:
            continue
        units = n.variables[field].units
        units =  " * ".join(map(regex_parser, units.split())) 
        data[field] = (np.squeeze(d), str(units))
        names[field] = n.variables[field].long_name.replace("_", " ")
        
    return bbox, data, names, dims, sizes


# At present, geographic plots in *yt* work best with global datasets, so for this example, we're using the global shear wave anomaly model of Simmons et al. 2010, available from IRIS [here](https://ds.iris.edu/ds/products/emc-gypsum/). So let's load and flatten our netcdf file:

# In[37]:


fi = "GYPSUM_percent.nc" 
n = nc4.Dataset(get_data_path(fi))
bbox, data, names, dims, sizes = prep_netcdf(n)


# Our global model covers the full range of latitude and longitude, but we have to currently set our bounds just under before creating our *yt* dataset:

# In[38]:


bbox[1][0]=-89.9
bbox[1][1]=89.9
bbox[2][0]=-189.9
bbox[2][1]=189.9
ds = yt.load_uniform_grid(data, sizes, 1.0, geometry=("internal_geographic", dims),bbox=bbox)


# Now that we have our dataset loaded, we can make a slice plot, setting our map projection using the `set_mpl_projection` function:

# In[20]:


target_depth = 100
center = ds.domain_center

p = yt.SlicePlot(ds,"depth",'dvs',center=[target_depth,center[1],center[2]]) # will be domain center
p.set_cmap('dvs', 'dusk_r')
p.set_zlim('dvs', -10, 10)
p.set_log('dvs', False)

p.annotate_title(f"Depth={target_depth} km")

p.set_mpl_projection(('PlateCarree', (), {'central_longitude':center[2]}))#{'central_longitude':-100, 'central_latitude':40}))
p._setup_plots()

p.plots['dvs'].axes.coastlines()


# In[21]:


p.show()


# We can easily switch projections, in this case to a Mollweide projection:

# In[22]:


p = yt.SlicePlot(ds,"depth",'dvs',center=[target_depth,center[1],center[2]]) # will be domain center
p.set_cmap('dvs', 'dusk_r')
p.set_zlim('dvs', -10, 10)
p.set_log('dvs', False)

p.annotate_title(f"Depth={target_depth} km")

p.set_mpl_projection(('Mollweide', (), {'central_longitude':center[2]}))#{'central_longitude':-100, 'central_latitude':40}))
p._setup_plots()

p.plots['dvs'].axes.coastlines()


# In[23]:


p.show()


# ### adding annotations
# 
# *yt* allows access to the underlying Cartopy axes, so we can easily add data annotations directly to the axes using standard matplotlib/cartopy commands. 
# 
# First, let's pull some data. Here, we're using obspy to query the IRIS catalogue for information about all earthquakes with a minimum magnitude of 5 from 2019-01-01 through 2020-11-01: 

# In[42]:


from obspy import UTCDateTime
from obspy.clients.fdsn import Client
client = Client("IRIS")
starttime = UTCDateTime("2019-01-01")
endtime = UTCDateTime("2020-11-01")
cat = client.get_events(starttime=starttime, endtime=endtime,minmagnitude=5)


# Now we can select the event locations:

# In[44]:


lat_lon_dep = [[ev.origins[0].latitude,ev.origins[0].longitude,ev.origins[0].depth] for ev in cat.events]
lat_lon_dep = np.array(lat_lon_dep)
lat_lon_dep.shape


# and now add them to our plot by first calling `_setup_plots` and then accesing the `plots` dictionary:

# In[45]:


p = yt.SlicePlot(ds,"depth",'dvs',center=[target_depth,center[1],center[2]]) # will be domain center
p.set_cmap('dvs', 'dusk_r')
p.set_zlim('dvs', -10, 10)
p.set_log('dvs', False)

p.annotate_title(f"Depth={target_depth} km")


p.set_mpl_projection(('Mollweide', (), {'central_longitude':center[2]}))#{'central_longitude':-100, 'central_latitude':40}))
p._setup_plots()


p.plots['dvs'].axes.scatter(lat_lon_dep[:,1],lat_lon_dep[:,0],
           transform=ccrs.PlateCarree(),color=[.8,.8,.8,.5])

p.plots['dvs'].axes.coastlines()
p.show()


# we can also use Cartopy's built in shapefile reader to add data. 
# 
# In the following cell, we read the Harvard Global Volcanic Database shapefile (link) containing locations of Holocene volcanos:

# In[46]:


from cartopy.io.shapereader import Reader

volc_file = get_data_path('harvard-glb-volc-shapefile/GLB_VOLC.shp')
volc_reader = Reader(volc_file)


# which we can again add to our plot:

# In[47]:


p = yt.SlicePlot(ds,"depth",'dvs',center=[target_depth,center[1],center[2]]) # will be domain center
p.set_cmap('dvs', 'dusk_r')
p.set_zlim('dvs', -10, 10)
p.set_log('dvs', False)

p.annotate_title(f"Depth={target_depth} km")


p.set_mpl_projection(('Mollweide', (), {'central_longitude':center[2]}))#{'central_longitude':-100, 'central_latitude':40}))
p._setup_plots()

points = list(volc_reader.geometries())
p.plots['dvs'].axes.scatter([point.x for point in points],
           [point.y for point in points],
           transform=ccrs.PlateCarree(),color='k')
    
p.plots['dvs'].axes.scatter(lat_lon_dep[:,1],lat_lon_dep[:,0],
           transform=ccrs.PlateCarree(),color=[.8,.8,.8,.5])

p.plots['dvs'].axes.coastlines()
p.show()


# In[48]:


p.save('./figures/seismic_fixed_depth.png')


# In[ ]: