<jupyter_start><jupyter_text>Load Data and Saved Model<jupyter_code>data_dir = '/home/teisberg/data/common_data/'
nc_measure = Dataset(os.path.join(data_dir, 'antarctica_ice_velocity_450m_v2.nc'))
nc_bedmachine = Dataset(os.path.join(data_dir, 'BedMachineAntarctica_2019-11-05_v01.nc'))
nc_bedmap2 = Dataset(os.path.join(data_dir, 'bedmap2_1km_input.nc'))

start_point = [550e3, -0.8e6] # Byrd
X_ctr, Y_ctr = start_point
size_x, size_y = 200e3, 200e3

xm = np.array(nc_measure.variables['x'])
ym = np.array(nc_measure.variables['y'])
mask_x = (xm > (X_ctr - size_x)) & (xm < (X_ctr + size_x))
mask_y = (ym > (Y_ctr - size_y)) & (ym < (Y_ctr + size_y))
measure_X, measure_Y = np.meshgrid(xm[mask_x], ym[mask_y])
measure_vx = np.array(nc_measure.variables['VX'])[mask_y, :][:, mask_x]
measure_vy = np.array(nc_measure.variables['VY'])[mask_y, :][:, mask_x]
api = wandb.Api()
run = api.run("teisberg/igarss2021/1p369fae")
for file in run.files():
    if file.name.startswith('gen_model/'):
        file.download(root='downloaded_model/', replace=True)
    if file.name == 'config.yaml':
        file.download(root='downloaded_model/', replace=True)
model_path = 'downloaded_model/gen_model'
parameter_yaml_filename = 'downloaded_model/config.yaml'

#
# Parameters
#

with open(parameter_yaml_filename) as f:
    config = yaml.load(f, Loader=yaml.FullLoader)
    
for k in config:
    if isinstance(config[k], dict):
        config[k] = config[k]['value']

config['wandb'] = False
print(config)

#
# Data Loading
#

with open(config['input_data_filename'], 'rb') as f:
    data = pickle.load(f)

dataset = PINNDataset(data, batch_size=config['batch_size'],
                        mode=config['mode'], n_random=config['n_random_points'])

#
# Create model
#

model = PINN(config, dataset, gen_model_filename=model_path)<jupyter_output><empty_output><jupyter_text>Whole Domain<jupyter_code>visualizer_callback = GeneratorVisualizationCallback(config, dataset, model)

roi = {
    'x0': 410000,
    'x1': 560000,
    'y1': -780000,
    'y0': -930000,
    'spacing': 500
}

x, y, pred = visualizer_callback.roi_prediction(roi)
h = np.reshape(np.array(pred['h']), np.shape(x))
vx = np.reshape(np.array(pred['vx']), np.shape(x))
vy = np.reshape(np.array(pred['vy']), np.shape(x))

dhdx = np.diff(h, axis=0)[:,:-1] / roi['spacing']
dhdy = np.diff(h, axis=1)[:-1,:] / roi['spacing']
dvxdx = np.diff(vx, axis=0)[:,:-1] / roi['spacing']
dvydy = np.diff(vy, axis=1)[:-1,:] / roi['spacing']
fig, axs = plt.subplots(2, 2, figsize=(6,5), facecolor='white', sharex=True, sharey=True)

h_norm = matplotlib.colors.Normalize(vmin=0, vmax=4000)

pcm0 = axs[0,0].scatter(dataset.r['x']/1000, dataset.r['y']/1000, c=dataset.r['h'], norm=h_norm, s=0.05, rasterized=True)
clb0 = fig.colorbar(pcm0, ax=axs[0,0])
clb0.set_label('Radar Ice Thickness [m]')
axs[0,0].set_aspect('equal')
axs[0,0].set_ylabel('Y [km]')

pcm1 = axs[1,0].pcolormesh(x/1000,y/1000,h, norm=h_norm, rasterized=True, shading='nearest')
clb1 = fig.colorbar(pcm1, ax=axs[1,0])
clb1.set_label('Interpolated Ice Thickness [m]')
axs[1,0].set_aspect('equal')
axs[1,0].set_xlabel('X [km]')
axs[1,0].set_ylabel('Y [km]')

pcm2 = axs[0,1].pcolormesh(measure_X/1000, measure_Y/1000, np.sqrt(measure_vx**2 + measure_vy**2), rasterized=True, shading='nearest')
arrow_subsample = 20
axs[0,1].quiver(measure_X[::arrow_subsample,::arrow_subsample]/1000,
              measure_Y[::arrow_subsample,::arrow_subsample]/1000,
              measure_vx[::arrow_subsample,::arrow_subsample]/1000,
              measure_vy[::arrow_subsample,::arrow_subsample]/1000,
              angles='xy', color='white', scale_units='xy', scale=5e-2, rasterized=True)
clb2 = fig.colorbar(pcm2, ax=axs[0,1])
clb2.set_label('MEaSURE Ice Velocity [m/yr]')
axs[0,1].set_aspect('equal')

axs[0,1].tick_params(axis='x', labelrotation=45) 
axs[1,0].tick_params(axis='x', labelrotation=45)
axs[0,0].tick_params(axis='x', labelrotation=45)

axs[0,0].set_xlim(roi['x0']/1000, roi['x1']/1000)
axs[0,0].set_ylim(roi['y0']/1000, roi['y1']/1000)

fig.tight_layout()
fig.savefig('figures/2d-results.svg', format='svg', dpi=1000)
# Context plot
crs_3031 = cartopy.crs.Stereographic(central_latitude=-90, true_scale_latitude=-71)
fig, ax = plt.subplots(figsize=(3,4), facecolor='white', subplot_kw=dict(projection=crs_3031)) # Create our plot

ax.set_xlim(-2.8e6,2.8e6)
ax.set_ylim(-2.8e6,2.8e6)
ax.coastlines(resolution='110m', color='black') # Add coastlines

rect = matplotlib.patches.Rectangle((roi['x0'], roi['y0']),
                                    roi['x1']-roi['x0'], roi['y1']-roi['y0'],
                                    linewidth=1, edgecolor='red', facecolor='none', linestyle='--')
ax.add_patch(rect)
fig.savefig('figures/context-map.svg', format='svg', dpi=1000)
def bedmachine_region(roi, field='thickness'):
    x = np.array(nc_bedmachine.variables['x'])
    y = np.array(nc_bedmachine.variables['y'])
    mask_x = (x >= roi['x0']) & (x <= roi['x1'])
    mask_y = (y >= roi['y0']) & (y <= roi['y1'])
    bm_X, bm_Y = np.meshgrid(x[mask_x], y[mask_y])
    bm_h = np.array(nc_bedmachine.variables[field])[mask_y, :][:, mask_x]
    #bm_source = np.array(nc_bedmachine.variables['source'])[mask_y, :][:, mask_x]
    #bm_mask = np.array(nc_bedmachine.variables['mask'])[mask_y, :][:, mask_x]
    return bm_X, bm_Y, bm_h

def bedmap2_region(roi):
    x = np.array(nc_bedmap2.variables['x'])
    y = np.array(nc_bedmap2.variables['y'])
    mask_x = (x >= roi['x0']) & (x <= roi['x1'])
    mask_y = (y >= roi['y0']) & (y <= roi['y1'])
    bm_X, bm_Y = np.meshgrid(x[mask_x], y[mask_y])
    bm_h = np.array(nc_bedmap2.variables['thk'])[mask_y, :][:, mask_x]
    #bm_source = np.array(nc_bedmachine.variables['source'])[mask_y, :][:, mask_x]
    #bm_mask = np.array(nc_bedmachine.variables['mask'])[mask_y, :][:, mask_x]
    return bm_X, bm_Y, bm_h
fig, axs = plt.subplots(1,3, figsize=(6,2), sharex=True, sharey=True)

h_norm = matplotlib.colors.Normalize(vmin=0, vmax=3200)

pcm0 = axs[0].pcolormesh(x/1000,y/1000,h, norm=h_norm, rasterized=True, shading='nearest')
clb0 = fig.colorbar(pcm0, ax=axs[0])
axs[0].set_aspect('equal')
axs[0].set_title('Ours')
axs[0].set_ylabel('Y [km]')

bm_X, bm_Y, bm_H = bedmachine_region(roi)
pcm1 = axs[1].pcolormesh(bm_X/1000,bm_Y/1000,bm_H, norm=h_norm, rasterized=True, shading='nearest')
clb1 = fig.colorbar(pcm1, ax=axs[1])
axs[1].set_aspect('equal')
axs[1].set_title('BedMachine Antarctica')

bm_X, bm_Y, bm_H = bedmap2_region(roi)
pcm2 = axs[2].pcolormesh(bm_X/1000,bm_Y/1000,bm_H, norm=h_norm, rasterized=True, shading='nearest')
clb2 = fig.colorbar(pcm2, ax=axs[2])
clb2.set_label('Ice Thickness [m]')
axs[2].set_aspect('equal')
axs[2].set_title('BedMap2')

for i in range(3):
    axs[i].tick_params(axis='x', labelrotation=45) 
    axs[i].set_xlabel('X [km]')
    
fig.savefig('figures/bm-comparison.svg', format='svg', dpi=1000)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19953663481030987, 'fraction_numerical': 0.04821894005212858, 'mean_word_length': 5.485446009389672, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '589331', 'license': 'MIT', 'n_tokens_mistral': 3139, 'n_tokens_neox': 2982, 'n_words': 407, 'path': 'figures-2d.ipynb', 'repo_name': 'thomasteisberg/igarss2021', 'size': '6906'}