From 151a69cf5c4949ec655bd0386fc1ba60b2a21d8e Mon Sep 17 00:00:00 2001 From: Mansa Krishna <73682539+mansakrishna23@users.noreply.github.com> Date: Sat, 6 Jul 2024 13:06:16 -0400 Subject: [PATCH] Revert "surf slopes added + combined plotting functions" --- pinnicle/utils/plotting.py | 217 +++++++++++++++++++++++++++++-------- tests/test_pinn.py | 55 ++++++++++ 2 files changed, 224 insertions(+), 48 deletions(-) diff --git a/pinnicle/utils/plotting.py b/pinnicle/utils/plotting.py index 0463570..5849231 100644 --- a/pinnicle/utils/plotting.py +++ b/pinnicle/utils/plotting.py @@ -227,6 +227,139 @@ 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 """ @@ -283,78 +416,66 @@ def plot_residuals(pinn, cmap='RdBu', cbar_bins=10, cbar_limits=[-5e3, 5e3]): return fig, axs -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 +def tripcolor_similarity(pinn, feature_name, feat_title=None, sim='MAE', cmap='jet', scale=1, colorbar_bins=10): + """tripcolor similarity, plot with ISSM triangulation """ - if feat_title=None: - if type(feature)==list: - raise TypeError('feat_title must be provided as input str') + if feat_title == None: + if type(feature_name)==list: + raise TypeError('feat_title must be provided as a str type input.') else: - feat_title=feature + feat_title = feature_name - fig, axs = plt.subplots(1, 3, figsize=figsize) + # initialize figure + fig, axs = plt.subplots(1, 3, figsize=(15, 4)) - # neural network inputs and outputs - ins = pinn.nn.parameters.input_variables - outs = pinn.nn.parameters.output_variables + # neural network features + input_names = pinn.nn.parameters.input_variables + output_names = pinn.nn.parameters.output_variables # inputs - 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])) + 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 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) + # 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])) - # grabbing the feature + # triangulation from ISSM (matlab --> python indexing) + elements = pinn.model_data.data['ISSM'].mesh_dict['elements']-1 + triangles = mpl.tri.Triangulation(meshx, meshy, elements) + + # grab feature + # initializing ref and pred 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)==list: - for feat in feature: - fid = outs.index(feat) + 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 = outs.index(feature) + fid = output_names.index(feature_name) ref_sol = np.squeeze(sol[:, fid:fid+1]*scale) pred_sol = np.squeeze(pred[:, fid:fid+1]*scale) - 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 + [cmin, cmax] = [0.9*np.min(np.append(ref_sol, pred_sol)), 1.1*np.max(np.append(ref_sol, pred_sol))] 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 plot + # looping through the columns of the 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) @@ -363,7 +484,7 @@ def plot_similarity(pinn, feature, feat_title=None, mdata='ISSM', figsize=(15, 4 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() == 'RSME': + elif sim.upper() == 'RMSE': diff = (ref_sol-pred_sol)**2 diff_val = np.round(np.sqrt(np.mean(diff)), 2) diff = np.sqrt(diff) @@ -385,7 +506,7 @@ def plot_similarity(pinn, feature, feat_title=None, mdata='ISSM', figsize=(15, 4 axes = axs[c].tripcolor(triangles, data_list[c], cmap=cmap, norm=colors.Normalize(vmin=cmin, vmax=cmax)) title = title_list[c] - # common plot settings + # common settings axs[c].set_title(title, fontsize=14) axs[c].axis('off') cb = plt.colorbar(axes, ax=axs[c]) diff --git a/tests/test_pinn.py b/tests/test_pinn.py index 2cc1641..b4f4b35 100644 --- a/tests/test_pinn.py +++ b/tests/test_pinn.py @@ -217,6 +217,32 @@ 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 @@ -247,6 +273,35 @@ 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)