Merge pull request #40 from mansakrishna23/main

flexible function resplot for pre residuals

pinnicle/utils/__init__.py

@@ -1,4 +1,4 @@
 from .helper import *
 from .history import History
 from .data_misfit import get
-from .plotting import plot_solutions, plot_dict_data, plot_data, plot_nn, plot_similarity, plot_residuals, tripcolor_similarity, tripcolor_residuals, diffplot
+from .plotting import plot_solutions, plot_dict_data, plot_data, plot_nn, tripcolor_similarity, tripcolor_residuals, diffplot, resplot

pinnicle/utils/plotting.py

@@ -347,188 +347,68 @@ def diffplot(pinn, feature, feat_title=None, mdata='ISSM', sim='mae', figsize=(1
 
     return fig, 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
+def resplot(pinn, mdata='ISSM', figsize=None, cmap='RdBu', cbar_bins=10, cbar_limits=[-5e3, 5e3], elements=None):
+    """plotting the pde residuals
     """
-    # 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]
+    X_ref = pinn.model_data.data[mdata].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())
+    # triangles / elements
+    if elements==None:
+        if pinn.model_data.data[mdata].mesh_dict == {}:
+            triangles = mpl.tri.Triangulation(meshx, meshy)
         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
-    """
-    input_names = pinn.nn.parameters.input_variables
-    output_names = pinn.nn.parameters.output_variables
+            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)
+    else:
+        triangles = elements
+        if len(triangles) != len(meshx):
+            raise ValueError('number of elements must equal number of (x, y) inputs')
 
-    # 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])
-
     Nr = len(pinn.physics.residuals)
     fig, axs = plt.subplots(1, len(pinn.physics.residuals), figsize=(5*Nr, 4))
-    levels = np.linspace(cbar_limits[0], cbar_limits[-1], 500)
-    # counting the pde residuals
+
     pde_dict = {} # counting the number of residuals per pde
     for i in pinn.params.physics.equations.keys():
         pde_dict[i] = 0
 
     for r in range(Nr):
         # looping through the equation keys
         for p in pinn.params.physics.equations.keys():
-            # check if the equation key is in the residual name
             if p in pinn.physics.residuals[r]:
                 pde_dict[p] += 1
                 pde_pred = pinn.model.predict(xref, operator=pinn.physics.operator(p))
-                op_pred = pde_pred[pde_dict[p]-1] # operator predicton
 
+                op_pred = pde_pred[pde_dict[p]-1]
                 if Nr <= 1:
-                    axes = axs.tricontourf(meshx, meshy, np.squeeze(op_pred), levels=levels, cmap=cmap, norm=colors.CenteredNorm())
+                    axes = axs.tripcolor(triangles, np.squeeze(op_pred), cmap=cmap, norm=colors.CenteredNorm(clip=[cbar_limits[0], cbar_limits[-1]]))
                     cb = plt.colorbar(axes, ax=axs)
                     cb.ax.tick_params(labelsize=14)
                     # adjusting the number of ticks
-                    colorbar_bins = ticker.MaxNLocator(nbins=cbar_bins)
-                    cb.locator = colorbar_bins
+                    num_bins = ticker.MaxNLocator(nbins=cbar_bins)
+                    cb.locator = num_bins
                     cb.update_ticks()
                     # setting the title
                     axs.set_title(str(pinn.physics.residuals[r]), fontsize=14)
                     axs.axis('off')
                 else:
-                    axes = axs[r].tricontourf(meshx, meshy, np.squeeze(op_pred), levels=levels, cmap=cmap, norm=colors.CenteredNorm())
+                    axes = axs[r].tripcolor(triangles, np.squeeze(op_pred), cmap=cmap, norm=colors.CenteredNorm(clip=[cbar_limits[0], cbar_limits[-1]]))
                     cb = plt.colorbar(axes, ax=axs[r])
                     cb.ax.tick_params(labelsize=14)
                     # adjusting the number of ticks
-                    colorbar_bins = ticker.MaxNLocator(nbins=cbar_bins)
-                    cb.locator = colorbar_bins
+                    num_bins = ticker.MaxNLocator(nbins=cbar_bins)
+                    cb.locator = num_bins
                     cb.update_ticks()
                     # title
                     axs[r].set_title(str(pinn.physics.residuals[r]), fontsize=14)

tests/test_plotting.py

@@ -2,7 +2,7 @@
 import pinnicle as pinn
 import numpy as np
 import deepxde as dde
-from pinnicle.utils import plot_similarity, plot_residuals, tripcolor_similarity, tripcolor_residuals, diffplot
+from pinnicle.utils import tripcolor_similarity, tripcolor_residuals, diffplot, resplot
 import matplotlib.pyplot as plt
 import pytest
 
@@ -50,34 +50,8 @@
 SSA["scalar_variables"] = {"B":1.26802073401e+08}
 hp["equations"] = {"SSA":SSA}
 
-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'])
-    plt.close("all") 
 
-def test_residuals(tmp_path):
+def test_resplot(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}
@@ -86,15 +60,15 @@ def test_residuals(tmp_path):
     experiment.compile()
     # plot_residuals(pinn, cmap='RdBu', cbar_bins=10, cbar_limits=[-5e3, 5e3])
     # default
-    fig, axs = plot_residuals(experiment)
+    fig, axs = resplot(experiment)
     assert (fig is not None) and (np.size(axs)==2)
-    fig, axs = plot_residuals(experiment, cmap='jet')
+    fig, axs = resplot(experiment, cmap='jet')
     assert (fig is not None) and (np.size(axs)==2)
-    fig, axs = plot_residuals(experiment, cbar_bins=5)
+    fig, axs = resplot(experiment, cbar_bins=5)
     assert (fig is not None) and (np.size(axs)==2)
-    fig, axs = plot_residuals(experiment, cbar_limits=[-1e4, 1e4])
+    fig, axs = resplot(experiment, cbar_limits=[-1e4, 1e4])
     assert (fig is not None) and (np.size(axs)==2)
-    fig, axs = plot_residuals(experiment, cmap='rainbow', cbar_bins=20, cbar_limits=[-7.5e3, 7.5e3])
+    fig, axs = resplot(experiment, cmap='rainbow', cbar_bins=20, cbar_limits=[-7.5e3, 7.5e3])
     assert (fig is not None) and (np.size(axs)==2)
 
     # add more physics, test again
@@ -104,7 +78,7 @@ def test_residuals(tmp_path):
     experiment = pinn.PINN(params=hp)
     experiment.compile()
 
-    fig, axs = plot_residuals(experiment)
+    fig, axs = resplot(experiment)
     assert (fig is not None) and (np.size(axs)==3)
     plt.close("all") 
 
@@ -157,6 +131,19 @@ def test_triresiduals(tmp_path):
     assert (fig is not None) and (np.size(axs)==2)
     plt.close("all") 
 
+def test_resplot(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()
+
+    fig, axs = resplot(experiment)
+    assert (fig is not None) and (np.size(axs)==2)
+    plt.close("all")
+
 def test_diffplot(tmp_path):
     hp["save_path"] = str(tmp_path)
     hp["is_save"] = True