1#!/usr/bin/env python
2# ----------------------------------------------------------------
3# Programmer(s):  Daniel R. Reynolds @ SMU
4# ----------------------------------------------------------------
5# SUNDIALS Copyright Start
6# Copyright (c) 2002-2020, Lawrence Livermore National Security
7# and Southern Methodist University.
8# All rights reserved.
9#
10# See the top-level LICENSE and NOTICE files for details.
11#
12# SPDX-License-Identifier: BSD-3-Clause
13# SUNDIALS Copyright End
14# ----------------------------------------------------------------
15# matplotlib-based plotting script for heat2D.cpp example
16# ----------------------------------------------------------------
17
18# imports
19import sys
20import numpy as np
21from pylab import *
22from mpl_toolkits.mplot3d import Axes3D
23from matplotlib import cm
24import matplotlib.pyplot as plt
25
26# determine the number of MPI processes used
27nprocs=1
28for i in range(1000):
29    sname = 'heat2d_subdomain.' + repr(i).zfill(3) + '.txt'
30    try:
31        f = open(sname,'r')
32        f.close()
33    except IOError:
34        nprocs = i
35        break
36
37# load subdomain information, store in table
38subdomains = np.zeros((nprocs,4), dtype=np.int)
39for i in range(nprocs):
40    sname = 'heat2d_subdomain.' + repr(i).zfill(3) + '.txt'
41    subd = np.loadtxt(sname, dtype=np.int)
42    if (i == 0):
43        nx = subd[0]
44        ny = subd[1]
45    else:
46        if ((subd[0] != nx) or (subd[1] != ny)):
47            sys.exit("error: subdomain files incompatible (clean up and re-run test)")
48    subdomains[i,:] = subd[2:6]
49
50# load first processor's data, and determine total number of time steps
51data = np.loadtxt('heat2d.000.txt', dtype=np.double)
52nt = np.shape(data)[0]
53
54# create empty array for all solution data
55results = np.zeros((nt,ny,nx))
56
57# insert first processor's data into results array
58istart = subdomains[0,0]
59iend = subdomains[0,1]
60jstart = subdomains[0,2]
61jend = subdomains[0,3]
62nxl = iend-istart+1
63nyl = jend-jstart+1
64for i in range(nt):
65    results[i,jstart:jend+1,istart:iend+1] = np.reshape(data[i,:], (nyl,nxl))
66
67# iterate over remaining data files, inserting into output
68if (nprocs > 1):
69    for isub in range(1,nprocs):
70        data = np.loadtxt('heat2d.' + repr(isub).zfill(3) + '.txt', dtype=np.double)
71        # check that subdomain has correct number of time steps
72        if (np.shape(data)[0] != nt):
73            sys.exit('error: subdomain ' + isub + ' has an incorrect number of time steps')
74        istart = subdomains[isub,0]
75        iend = subdomains[isub,1]
76        jstart = subdomains[isub,2]
77        jend = subdomains[isub,3]
78        nxl = iend-istart+1
79        nyl = jend-jstart+1
80        for i in range(nt):
81            results[i,jstart:jend+1,istart:iend+1] = np.reshape(data[i,:], (nyl,nxl))
82
83# determine extents of plots
84maxtemp = 1.1*results.max()
85mintemp = 0.9*results.min()
86
87# generate plots of results
88kx = 0.5
89ky = 0.75
90kprod = (kx+4.0*ky)*np.pi**2
91dt = 0.015
92for tstep in range(nt):
93
94    # set string constants for output plots, current time, mesh size
95    pname = 'heat2d_surf.' + repr(tstep).zfill(3) + '.png'
96    cname = 'heat2d_err.' + repr(tstep).zfill(3) + '.png'
97    tstr  = repr(tstep)
98    nxstr = repr(nx)
99    nystr = repr(ny)
100
101    # set x and y meshgrid objects
102    xspan = np.linspace(0.0, 1.0, nx)
103    yspan = np.linspace(0.0, 1.0, ny)
104    X,Y = np.meshgrid(xspan,yspan)
105
106    # plot current solution as a surface, and save to disk
107    fig = plt.figure(1)
108    ax = fig.add_subplot(111, projection='3d')
109    ax.plot_surface(X, Y, results[tstep,:,:], rstride=1, cstride=1,
110                    cmap=cm.jet, linewidth=0, antialiased=True, shade=True)
111    ax.set_xlabel('x')
112    ax.set_ylabel('y')
113    ax.set_zlim((mintemp, maxtemp))
114    ax.view_init(20,45)
115    title('u(x,y) at output ' + tstr + ', mesh = ' + nxstr + 'x' + nystr)
116    savefig(pname)
117    plt.close()
118
119    # plot error in current solution (as a contour, and save to disk)
120    t = tstep*dt;
121    at = (1.0 - np.exp(-t*kprod))/kprod
122    utrue = at*np.sin(np.pi*X)*np.sin(2.0*np.pi*Y);
123    uerr = np.abs(utrue - results[tstep,:,:])
124    plt.contourf(xspan,yspan,uerr,15, cmap=plt.cm.jet)
125    plt.colorbar()
126    plt.xlabel('x')
127    plt.ylabel('y')
128    plt.title('Error at output ' + tstr + ', mesh = ' + nxstr + 'x' + nystr)
129    plt.savefig(cname)
130    plt.close()
131
132
133
134
135##### end of script #####
136