Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

surf slopes added + combined plotting functions #35

Merged
merged 3 commits into from
Jul 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
217 changes: 48 additions & 169 deletions pinnicle/utils/plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,139 +227,6 @@ def plot_data(X, Y, im_data, axs=None, vranges={}, **kwargs):

return axs

def plot_similarity(pinn, feature_name, feat_title=None, sim='MAE', cmap='jet', scale=1, cols=[0, 1, 2], cbar_bins=10):
"""
plotting the similarity between reference and predicted
solutions, mae default

Args:
pinn : pinnicle.model
the trained PINN model
feature_name : str
the name of a predicted feature of the PINN.
for the L-2 norm of two or more predictions, write as a list. e.g., ['u', 'v'].
feat_title : str (default=None, will be set to feature_name if None provided)
the name of the predicted feature in the title.
sim : str (default='MAE')
the similarity/comparison type.
options include: 'MAE', 'MSE', 'RMSE', 'SIMPLE'
(can be written as upper case or lower case) e.g., 'MSE' and 'mse' will give the same result.
cmap : str (default='jet', for similarity default='RdBu')
the matplotlib colormap name as a str.
scale : float (default=1)
the scale by which to multiply predictions (e.g., m/s * yts = m/year, then scale = yts)
cols : list (default=[0, 1, 2])
can specify which columns of the figure to extract. 0 = reference, 1 = prediction, 2 = similarity.
cbar_bins : int (default=10)
the number of bins/ticks on the c-axis.

Returns:
fig, axs
plot of the reference, prediction, and similarity
"""
# setting the figure title
if feat_title == None:
if type(feature_name) == list:
raise TypeError('feat_title must be provided as an input string')
else:
feat_title = feature_name

# initialize figure, default all 3 columns
if len(cols) == 1:
# subplots returns a single Axes if only 1 figure, but we need array later
fig, ax_single = plt.subplots(1, len(cols), figsize=(5*len(cols), 4))
axs = [ax_single]
else:
fig, axs = plt.subplots(1, len(cols), figsize=(5*len(cols), 4))

# inputs and outputs of NN
input_names = pinn.nn.parameters.input_variables
output_names = pinn.nn.parameters.output_variables

# inputs
X_ref = pinn.model_data.data['ISSM'].X_dict
xref = X_ref[input_names[0]].flatten()[:,None]
for i in range(1, len(input_names)):
xref = np.hstack((xref, X_ref[input_names[i]].flatten()[:,None]))
meshx = np.squeeze(xref[:, 0])
meshy = np.squeeze(xref[:, 1])

# predictions
pred = pinn.model.predict(xref)

# reference solution
X_sol = pinn.model_data.data['ISSM'].data_dict
sol = X_sol[output_names[0]].flatten()[:,None] # initializing array
for i in range(1, len(output_names)):
sol = np.hstack((sol, X_sol[output_names[i]].flatten()[:,None]))

# grab feature
# initializing reference and prediction
ref_sol = np.zeros_like(np.squeeze(sol[:, 0:1]*scale))
pred_sol = np.zeros_like(np.squeeze(pred[:, 0:1]*scale))

if type(feature_name) == list:
for feat in feature_name:
fid = output_names.index(feat)
ref_sol += (np.squeeze(sol[:, fid:fid+1]*scale))**2
pred_sol += (np.squeeze(pred[:, fid:fid+1]*scale))**2
ref_sol = np.sqrt(ref_sol)
pred_sol = np.sqrt(pred_sol)
else:
fid = output_names.index(feature_name)
ref_sol = np.squeeze(sol[:, fid:fid+1]*scale)
pred_sol = np.squeeze(pred[:, fid:fid+1]*scale)

[cmin, cmax] = [np.min(np.append(ref_sol, pred_sol)), np.max(np.append(ref_sol, pred_sol))]
levels = np.linspace(cmin*0.9, cmax*1.1, 500)
data_list = [ref_sol, pred_sol]
title_list = [ feat_title + r"$_{ref}$", feat_title + r"$_{pred}$"]

# plotting
for c, col in enumerate(cols):
if col == 2:
if sim.upper() == 'MAE':
diff = np.abs(ref_sol-pred_sol)
diff_val = np.round(np.mean(diff), 2)
title = r"|"+feat_title+r"$_{ref} - $"+feat_title+r"$_{pred}$|, MAE="+str(diff_val)
elif sim.upper() == 'MSE':
diff = (ref_sol-pred_sol)**2
diff_val = np.round(np.mean(diff), 2)
title = r"$($"+feat_title+r"$_{ref} - $"+feat_title+r"$_{pred})^2$, MSE="+str(diff_val)
elif sim.upper() == 'RMSE':
diff = (ref_sol-pred_sol)**2
diff_val = np.round(np.sqrt(np.mean(diff)), 2)
diff = np.sqrt(diff)
title = r"$(($"+feat_title+r"$_{ref} - $"+feat_title+r"$_{pred})^2)^{1/2}$, RMSE="+str(diff_val)
elif sim.upper() == 'SIMPLE':
diff = ref_sol-pred_sol
diff_val = np.round(np.mean(diff), 2)
title = feat_title+r"$_{ref} - $"+feat_title+r"$_{pred}$, AVG. DIFF="+str(diff_val)
else:
print('Default similarity MAE implemented.')
diff = np.abs(ref_sol-pred_sol)
diff_val = np.round(np.mean(diff), 2)
title = r"|"+feat_title+r"$_{ref} - $"+feat_title+r"$_{pred}$|, MAE="+str(diff_val)

levels = np.linspace(np.min(diff)*0.9, np.max(diff)*1.1, 500)
data = np.squeeze(diff)
axes = axs[c].tricontourf(meshx, meshy, data, levels=levels, cmap='RdBu', norm=colors.CenteredNorm())
else:
axes = axs[c].tricontourf(meshx, meshy, data_list[col], levels=levels, cmap=cmap)
title = title_list[col]

# common settings
axs[c].set_title(title, fontsize=14)
cb = plt.colorbar(axes, ax=axs[c])
cb.ax.tick_params(labelsize=14)
colorbar_bins = ticker.MaxNLocator(nbins=cbar_bins)
cb.locator = colorbar_bins
cb.update_ticks()
axs[c].axis('off')

# save figure to path as defined
return fig, axs

def plot_residuals(pinn, cmap='RdBu', cbar_bins=10, cbar_limits=[-5e3, 5e3]):
"""plotting the pde residuals
"""
Expand Down Expand Up @@ -416,66 +283,78 @@ def plot_residuals(pinn, cmap='RdBu', cbar_bins=10, cbar_limits=[-5e3, 5e3]):

return fig, axs

def tripcolor_similarity(pinn, feature_name, feat_title=None, sim='MAE', cmap='jet', scale=1, colorbar_bins=10):
"""tripcolor similarity, plot with ISSM triangulation
def plot_similarity(pinn, feature, feat_title=None, mdata='ISSM', figsize=(15, 4), sim='mae', cmap='jet', clim=None, scale=1, colorbar_bins=10, elements=None):
"""plotting function: reference sol, prediction, and difference
"""
if feat_title == None:
if type(feature_name)==list:
raise TypeError('feat_title must be provided as a str type input.')
if feat_title=None:
if type(feature)==list:
raise TypeError('feat_title must be provided as input str')
else:
feat_title = feature_name
feat_title=feature

# initialize figure
fig, axs = plt.subplots(1, 3, figsize=(15, 4))
fig, axs = plt.subplots(1, 3, figsize=figsize)

# neural network features
input_names = pinn.nn.parameters.input_variables
output_names = pinn.nn.parameters.output_variables
# neural network inputs and outputs
ins = pinn.nn.parameters.input_variables
outs = pinn.nn.parameters.output_variables

# inputs
X_ref = pinn.model_data.data['ISSM'].X_dict
xref = X_ref[input_names[0]].flatten()[:,None]
for i in range(1, len(input_names)):
xref = np.hstack((xref, X_ref[input_names[i]].flatten()[:,None]))
X_ref = pinn.model_data.data[mdata].X_dict
xref = X_ref[ins[0]].flatten()[:,None]
for i in range(1, len(ins)):
xref = np.hstack((xref, X_ref[ins[i]].flatten()[:,None]))
meshx = np.squeeze(xref[:, 0])
meshy = np.squeeze(xref[:, 1])

# predictions
pred = pinn.model.predict(xref)

# reference solution
X_sol = pinn.model_data.data['ISSM'].data_dict
sol = X_sol[output_names[0]].flatten()[:,None]
for i in range(1, len(output_names)):
sol = np.hstack((sol, X_sol[output_names[i]].flatten()[:,None]))

# triangulation from ISSM (matlab --> python indexing)
elements = pinn.model_data.data['ISSM'].mesh_dict['elements']-1
triangles = mpl.tri.Triangulation(meshx, meshy, elements)
# reference sol
X_sol = pinn.model_data.data[mdata].data_dict
sol = X_sol[outs[0]].flatten()[:,None]
for i in range(1, len(outs)):
sol = np.hstack((sol, X_sol[outs[i]].flatten()[:,None]))

# elements, if any
if elements!=None:
triangles = elements
if len(triangles)!=len(meshx):
raise ValueError('number of elements must equal number number of vertices (x, y)')
else:
if pinn.model_data.data[mdata].mesh_dict == {}:
triangles = mpl.tri.Triangulation(meshx, meshy)
else:
if pinn.params.param_dict['data'][mdata]['data_path'].endswith('mat'):
elements = pinn.model_data.data[mdata].mesh_dict['elements']-1
else:
elements = pinn.model_data.data[mdata].mesh_dict['elements']
triangles = mpl.tri.Triangulation(meshx, meshy, elements)

# grab feature
# initializing ref and pred
# grabbing the feature
ref_sol = np.zeros_like(np.squeeze(sol[:, 0:1]*scale))
pred_sol = np.zeros_like(np.squeeze(pred[:, 0:1]*scale))
if type(feature_name) == list:
for feat in feature_name:
fid = output_names.index(feat)
if type(feature)==list:
for feat in feature:
fid = outs.index(feat)
ref_sol += (np.squeeze(sol[:, fid:fid+1]*scale))**2
pred_sol += (np.squeeze(pred[:, fid:fid+1]*scale))**2
ref_sol = np.sqrt(ref_sol)
pred_sol = np.sqrt(pred_sol)
else:
fid = output_names.index(feature_name)
fid = outs.index(feature)
ref_sol = np.squeeze(sol[:, fid:fid+1]*scale)
pred_sol = np.squeeze(pred[:, fid:fid+1]*scale)

[cmin, cmax] = [0.9*np.min(np.append(ref_sol, pred_sol)), 1.1*np.max(np.append(ref_sol, pred_sol))]
if clim==None:
[cmin, cmax] = [0.9*np.min(np.append(ref_sol, pred_sol)), 1.1*np.max(np.append(ref_sol, pred_sol))]
else:
[cmin, cmax] = clim
data_list = [ref_sol, pred_sol]
title_list = [ feat_title + r"$_{ref}$", feat_title + r"$_{pred}$"]
title_list = [feat_title + r"$_{ref}$", feat_title + r"$_{pred}$"]

# looping through the columns of the plot
# looping through plot
for c in range(3):
if c == 2:
if c==2:
if sim.upper() == 'MAE':
diff = np.abs(ref_sol-pred_sol)
diff_val = np.round(np.mean(diff), 2)
Expand All @@ -484,7 +363,7 @@ def tripcolor_similarity(pinn, feature_name, feat_title=None, sim='MAE', cmap='j
diff = (ref_sol-pred_sol)**2
diff_val = np.round(np.mean(diff), 2)
title = r"$($"+feat_title+r"$_{ref} - $"+feat_title+r"$_{pred})^2$, MSE="+str(diff_val)
elif sim.upper() == 'RMSE':
elif sim.upper() == 'RSME':
diff = (ref_sol-pred_sol)**2
diff_val = np.round(np.sqrt(np.mean(diff)), 2)
diff = np.sqrt(diff)
Expand All @@ -506,7 +385,7 @@ def tripcolor_similarity(pinn, feature_name, feat_title=None, sim='MAE', cmap='j
axes = axs[c].tripcolor(triangles, data_list[c], cmap=cmap, norm=colors.Normalize(vmin=cmin, vmax=cmax))
title = title_list[c]

# common settings
# common plot settings
axs[c].set_title(title, fontsize=14)
axs[c].axis('off')
cb = plt.colorbar(axes, ax=axs[c])
Expand Down
55 changes: 0 additions & 55 deletions tests/test_pinn.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,32 +217,6 @@ def test_plot(tmp_path):
assert len(im_data) == 5
assert im_data['u'].shape == (10,10)

def test_similarity(tmp_path):
hp["save_path"] = str(tmp_path)
hp["is_save"] = False
issm["data_size"] = {"u":100, "v":100, "s":100, "H":100, "C":None}
hp["data"] = {"ISSM": issm}
experiment = pinn.PINN(params=hp)
experiment.compile()
# plot_similarity(pinn, feature_name, sim='MAE', cmap='jet', scale=1, cols=[0, 1, 2])
# default
fig, axs = plot_similarity(experiment, feature_name='s')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = plot_similarity(experiment, feature_name='H', cols=[0])
assert (fig is not None) and (np.size(axs) == 1)
fig, axs = plot_similarity(experiment, feature_name="u", sim="mae", cols=[2])
assert (fig is not None) and (np.size(axs) == 1)
fig, axs = plot_similarity(experiment, feature_name="v", sim="Mse", cols=[2, 1])
assert (fig is not None) and (np.size(axs) == 2)
fig, axs = plot_similarity(experiment, feature_name="C", sim="rmse", cols=[0, 2, 1])
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = plot_similarity(experiment, feature_name="H", sim="SIMPLE")
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = plot_similarity(experiment, feature_name=['u', 'v'], feat_title='vel')
assert (fig is not None) and (np.size(axs) == 3)
with pytest.raises(TypeError):
fig, axs = plot_similarity(experiment, feature_name=['u', 'v'])

def test_residuals(tmp_path):
hp["save_path"] = str(tmp_path)
hp["is_save"] = False
Expand Down Expand Up @@ -273,35 +247,6 @@ def test_residuals(tmp_path):
fig, axs = plot_residuals(experiment)
assert (fig is not None) and (np.size(axs)==3)

def test_trisimilarity(tmp_path):
hp["equations"] = {"SSA":SSA}
hp["save_path"] = str(tmp_path)
hp["is_save"] = False
issm["data_size"] = {"u":100, "v":100, "s":100, "H":100, "C":None}
hp["data"] = {"ISSM": issm}
experiment = pinn.PINN(params=hp)
experiment.compile()
# plot_similarity(pinn, feature_name, sim='MAE', cmap='jet', scale=1, cols=[0, 1, 2])
# default
fig, axs = tripcolor_similarity(experiment, feature_name='s')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name='s', sim='mae')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name='s', sim='SIMPLE')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name='s', cmap='terrain')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name='s', sim='Rmse')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name='s', sim='mse')
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name='s', colorbar_bins=5)
assert (fig is not None) and (np.size(axs) == 3)
fig, axs = tripcolor_similarity(experiment, feature_name=['u', 'v'], feat_title='vel', scale=experiment.model_data.yts)
assert (fig is not None) and (np.size(axs) == 3)
with pytest.raises(TypeError):
fig, axs = tripcolor_similarity(experiment, feature_name=['u', 'v'])

def test_triresiduals(tmp_path):
hp["equations"] = {"SSA":SSA}
hp["save_path"] = str(tmp_path)
Expand Down
Loading