<jupyter_start><jupyter_text>tgb - 5/29/2019 - Meant as a simplification of https://github.com/tbeucler/CBRAIN-CAM/blob/master/notebooks/stephans-devlog/6.2-Debug-cut-off-runs.ipynb.  The goals will be:  1) To calculate the Jacobian without eager execution following https://gist.github.com/raspstephan/96cba1283adcbeb3dab553b5494481e3file-jac-no-eager-ipynb  2) To map the min and max of the eigenvalues around the instability to show the decrease in range and increase in min as an instability develops  1) Load CBRAIN and 8-col case<jupyter_code>from cbrain.imports import *
from cbrain.data_generator import *
from cbrain.utils import *
from cbrain.normalization import *
%matplotlib inline
CAMDIR = '/scratch/05488/tg847872/revision_debug/'
cam_ds = xr.open_mfdataset(CAMDIR+ 'debug.cam2.h1.*.nc', decode_times=False, concat_dim='time')
in_vars = ['QBP', 'TBP', 'VBP', 'PS', 'SOLIN', 'SHFLX', 'LHFLX']
out_vars = ['PHQ', 'TPHYSTND', 'FSNT', 'FSNS', 'FLNT', 'FLNS', 'PRECT']

scale_dict = load_pickle('../../nn_config/scale_dicts/002_pnas_scaling.pkl')

input_transform = InputNormalizer(
    xr.open_dataset(f'{CAMDIR}/001_norm.nc'),
    in_vars,
    'mean', 'maxrs',
    var_cut_off={'QBP':14, 'TBP':14}
)
from cbrain.models import fc_model
model = fc_model(66, 65, [256]*9, 'LeakyReLU')
model.load_weights(
    f'{CAMDIR}weights.h5'
)<jupyter_output>WARNING:tensorflow:From /home1/05823/tg851228/miniconda3/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:435: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Colocations handled automatically by placer.<jupyter_text>2) Crash animation using Holoview<jupyter_code>import holoviews as hv
hv.extension('matplotlib')
def anim(da, dims=['lon', 'lat'], **kwargs):
    ds = hv.Dataset(da)
    return ds.to(hv.Image, dims, rtol=100).options(**kwargs)
cam_ds = xr.open_mfdataset(CAMDIR+ 'debug.cam2.h1.*.nc', 
                           decode_times=False, concat_dim='time')
cam_ds.PHQ.isel(time=-50, lev=-1).plot()<jupyter_output><empty_output><jupyter_text>3) Jacobian functions without eager_execution<jupyter_code>from tensorflow.python.ops.parallel_for.gradients import batch_jacobian
def get_jacobian(x, model):
    sess = tf.keras.backend.get_session()
    jac = jacobian(model.output, model.input)
    J = sess.run(jac, feed_dict={model.input: x.astype(np.float32)[None]})
    return J.squeeze()
def get_batch_jacobian(x, model):
    sess = tf.keras.backend.get_session()
    jac = batch_jacobian(model.output, model.input)
    J = sess.run(jac, feed_dict={model.input: x.astype(np.float32)})
    return J.squeeze()
from tensorflow.python.platform import flags
from numpy import linalg as LA

NN = model
t_array = [-50]
Nt = len(t_array)
TQ_sp = np.zeros((Nt,128,2))
Q_sp = np.copy(TQ_sp)
T_sp = np.copy(TQ_sp)
coormax = np.zeros((Nt,2))

cf_inp = np.zeros((1,66))
for index in range (66):
    if index<16: cf_inp[0,index]=L_V;
    elif index<32: cf_inp[0,index]=C_P;
    else: cf_inp[0,index]=1;
        
cf_oup = np.zeros((1,65))
for index in range (65):
    if index<30: cf_oup[0,index]=L_V;
    elif index<60: cf_oup[0,index]=C_P;
    else: cf_oup[0,index]=1;

for it,t in enumerate(t_array):
    ilat_max = np.argmax(cam_ds.PHQ.isel(time=t, lev=-1).max('lon').values)
    ilon_max = np.argmax(cam_ds.PHQ.isel(time=t, lev=-1).max('lat').values)
    print('it=',it,'/',Nt-1,' t=',t,' ilat,ilon,lat,lon=',\
          ilat_max, ilon_max, \
          cam_ds.lat[ilat_max].values, cam_ds.lon[ilon_max].values)
        
    for ilon in range(128):

        print('ilon=',ilon,'/127 ',end="\r")

        in_vec = np.concatenate([
            cam_ds.NNQBP.isel(time=t, lat=ilat_max, lon=ilon, lev=slice(14, None)).values,
            cam_ds.NNTBP.isel(time=t, lat=ilat_max, lon=ilon, lev=slice(14, None)).values,
            cam_ds.NNVBP.isel(time=t, lat=ilat_max, lon=ilon).values,
            cam_ds.NNPS.isel(time=t, lat=ilat_max, lon=ilon).values[None],
            cam_ds.NNSOLIN.isel(time=t, lat=ilat_max, lon=ilon).values[None],
            cam_ds.NNSHF.isel(time=t, lat=ilat_max, lon=ilon).values[None],
            cam_ds.NNLHF.isel(time=t, lat=ilat_max, lon=ilon).values[None],
        ])[None, :]

        inp = input_transform.transform(in_vec)
        J = get_jacobian(inp[0], NN)
model<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.15445859872611464, 'fraction_numerical': 0.03230209281164695, 'mean_word_length': 4.349148418491485, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 12, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '154109', 'license': 'MIT', 'n_tokens_mistral': 1849, 'n_tokens_neox': 1733, 'n_words': 307, 'path': 'notebooks/tbeucler_devlog/016_Stampede2_Jacobians.ipynb', 'repo_name': 'wallingTACC/CBRAIN-CAM', 'size': '4396'}