Merge pull request #16 from rezgarshakeri/leila/review

Plotting added

.gitignore

@@ -1,2 +1,4 @@
 *.pyc
 ssshtest
+*.png
+

.travis.yml

@@ -19,3 +19,4 @@ script:
     - pycodestyle src/test_FE_subroutines.py
     - pycodestyle src/heat2d.py
     - pycodestyle src/FE_subroutines.py
+    - pycodestyle src/plot.py    

README.md

@@ -7,10 +7,10 @@ This code can be run with:
 
 `python src/heat2d.py`
 
-If you want to change the number of elements in x and y directions, or change the value of boundary conditions you can run with:
+One can change the number of elements in x and y directions, or the value of boundary conditions with:
 
 ```
-python heat2d.py --num_elm_x 4 --num_elm_y 4 --T0_bottom 0 --T0_left -5 --heat_source 5 --flux_top 10
+python src/heat2d.py --num_elm_x 20 --num_elm_y 20 --T0_bottom 50 --T0_left -50 --heat_source 400 --flux_top 100 --grid
 ```
 
 `--num_elm_x`: number of element in x direction
@@ -25,4 +25,12 @@ python heat2d.py --num_elm_x 4 --num_elm_y 4 --T0_bottom 0 --T0_left -5 --heat_s
 
 `--flux_top`: heat flux on the top boundary
 
-The output of the above command is the solution array which is the `Temperature` of all nodes in the discrete domain.
+`--grid`: plotting the grid if the user specifies
+
+The output of the above command contains the solution array which is the `Temperature` of all nodes in the discrete domain and a contour plot created illustrating the temperature distribution. The created plots for the shown command, with and without grid, is as following.
+
+**Temperature distribution with the grid:**
+<center><img src="images/ex_grid.png" width="300"/></center>  
+
+**Temperature distribution without the grid:**
+<center><img src="images/ex_nogrid.png" width="300"/></center>

src/FE_subroutines.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import math
+import sys
 
 
 def setup(nelx, nely):
@@ -105,15 +106,17 @@ def connectivity(nelx, nely):
 
     # Total number of elements in the domain
     nel = nelx*nely
-    IEN = np.zeros((4, nel), dtype=int)
+    IEN = np.zeros((nen, nel), dtype=int)
     rowcount = 0
-    for elementcount in range(0, nel):
-        IEN[0][elementcount] = elementcount + rowcount
-        IEN[1][elementcount] = elementcount + 1 + rowcount
-        IEN[2][elementcount] = elementcount + (lpx + 1) + rowcount
-        IEN[3][elementcount] = elementcount + (lpx) + rowcount
-        if np.mod(elementcount + 1, lpx - 1) == 0:
-            rowcount = rowcount + 1
+    # Connectivity matrix for 4-node elements
+    if nen == 4:
+        for elementcount in range(0, nel):
+            IEN[0][elementcount] = elementcount + rowcount
+            IEN[1][elementcount] = elementcount + 1 + rowcount
+            IEN[2][elementcount] = elementcount + (lpx + 1) + rowcount
+            IEN[3][elementcount] = elementcount + (lpx) + rowcount
+            if np.mod(elementcount + 1, lpx - 1) == 0:
+                rowcount = rowcount + 1
 
     return IEN
 
@@ -150,11 +153,11 @@ def Dirichlet_BCs(nelx, nely, T0_bottom, T0_left):
 
     # Essential B.C. (prescribed temperature)
     for i in range(0, lpx):
-        flags[i] = 2
+        flags[i] = 1
         e_bc[i] = T0_bottom       # bottom edge
 
     for i in range(lpx, nnp - nelx, lpx):
-        flags[i] = 2
+        flags[i] = 1
         e_bc[i] = T0_left         # left edge
 
     return e_bc, flags
@@ -204,7 +207,7 @@ def setup_ID_LM(nelx, nely, T0_bottom, T0_left):
     count = 0
     count1 = 0
     for i in range(0, neq):
-        if flags[i] == 2:
+        if flags[i] == 1:
             ID[i][0] = count
             d0[count] = e_bc[i]
             count = count + 1
@@ -264,8 +267,8 @@ def d_basis(xi, eta, coord):
 
     """
     # Derivative of N
-    #   the first row with respect to xi
-    #   the second row with respect to eta
+    #   dN[0, :] = dN/dxi
+    #   dN[0, :] = dN/deta
     dN = 0.25*np.array([[eta-1, 1-eta, 1+eta, -eta-1],
                         [xi-1, -xi-1, 1+xi, 1-xi]])
 
@@ -317,7 +320,10 @@ def gauss(ngp):
         w = np.array([1, 1])
         return w, gp
     else:
-        print("Error: This code supports only 1 or 2 quadrature points.")
+        print("Error: This code supports only 2 quadrature points so far!")
+        sys.exit(1)
+
+    return None
 
 
 def heat2delem(nelx, nely, ngp, s0, e):

src/heat2d.py

@@ -4,6 +4,7 @@
 """
 
 import FE_subroutines as FE
+import plot
 import argparse
 
 
@@ -45,6 +46,9 @@ def main():
                         type=float,
                         default=0,
                         help='Prescribed flux on the top edge')
+    parser.add_argument('--grid',
+                        action="store_true",
+                        help='Option for plotting the grid')
     args = parser.parse_args()
 
     nelx = args.num_elm_x
@@ -53,6 +57,7 @@ def main():
     T0_left = args.T0_left
     s0 = args.heat_source
     flux_top = args.flux_top
+    grid = args.grid
 
     # Gauss point, 2 is better than 1 since it has less errors
     ngp = 2
@@ -72,6 +77,11 @@ def main():
     # Print the nodal Temperature
     print(d)
 
+    # Plot the 2d temperature distribution
+    plot.plot_temp(nelx, nely, T0_bottom, T0_left, K, F, grid)
+
+    return
+
 
 if __name__ == '__main__':
     main()

src/plot.py

@@ -0,0 +1,102 @@
+"""This file includes function for plotting
+the Temperature distribution in a 2d domain.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import FE_subroutines as fe
+
+
+def plot_temp(nelx, nely, T0_bottom, T0_left, K, F, grid):
+    """This function plots the Temperature distribution
+    in a 2d domain.
+
+    Input:
+    ------
+    nelx:       integer
+                number of elements in the x direction.
+    nely:       integer
+                number of elements in the y direction.
+    T0_bottom:  float
+                prescribed temperature on the bottom edge
+    T0_left:    float
+                prescribed temperature on the left edge
+    K:          float (2d array)
+                assembled stiffness matrix
+    F:          float (1d array)
+                assembled forcing vector
+    grid:       boolean
+                true, if the user chooses to plot the grid
+
+    Output:
+    -------
+    A contour plot for the temperature distriboution which
+    looks like the following sketch for nelx = 2, nely = 2
+    6---------7----------8
+    |         |   (3)    |
+    |   (2)   |      ----5
+    |      ---4-----/    |
+    3-----/   |   (1)    |
+    |         |      ----2
+    |   (0)   |     /
+    |     ----1----/
+    0----/
+    """
+
+    # Calls functions to return required output for the plots
+    d0, ID, LM = fe.setup_ID_LM(nelx, nely, T0_bottom, T0_left)
+    d = fe.solvedr(nelx, nely, K, F, d0)
+    nel, lpx, lpy, nnp, ndof, nen, neq = fe.setup(nelx, nely)
+    x, y = fe.phys_coord(nelx, nely)
+    IEN = fe.connectivity(nelx, nely)
+
+    # Creates space for a figure to be drawn
+    fig, ax = plt.subplots()
+
+    # Arrange the temperature array in the physical ordering
+    d1 = np.zeros(nnp)
+    for i in range(nnp):
+        d1[i] = d[ID[i][0]][0]
+
+    # Creat and plot elements/grid
+    if grid:
+        xx = np.zeros(nen + 1)
+        yy = np.zeros(nen + 1)
+        dd = np.zeros(nen + 1)
+        for i in range(nel):
+            for j in range(nen):
+                xx[j] = x[IEN[j, i]]
+                yy[j] = y[IEN[j, i]]
+                dd[j] = d1[IEN[j, i]]
+            xx[nen] = x[IEN[0, i]]
+            yy[nen] = y[IEN[0, i]]
+            dd[nen] = d1[IEN[0, i]]
+            plt.fill(xx, yy, edgecolor='black', fill=False)
+
+    # Create and populate 2D Grid and the node's corresponding temperature
+    X = np.zeros((lpx, lpy))
+    Y = np.zeros((lpx, lpy))
+    D = np.zeros((lpx, lpy))
+    for i in range(lpx):
+        for j in range(lpy):
+            X[i][j] = x[j*lpx + i][0]
+            Y[i][j] = y[j*lpx + i][0]
+            D[i][j] = d1[j*lpx + i]
+
+    # Turn off the right and top sides of the bounding box
+    right_side = ax.spines["top"]
+    right_side.set_visible(False)
+    top_side = ax.spines["right"]
+    top_side.set_visible(False)
+
+    # Plot
+    plt.contourf(X, Y, D, cmap=plt.get_cmap('coolwarm'))
+    plt.colorbar()
+    plt.title('Temperature distribution')
+    plt.xlabel('X')
+    plt.ylabel('Y')
+    plt.axis('equal')
+    fig.tight_layout()
+    plt.savefig('fig.png', bbox_inches='tight')
+
+    return

src/test_FE_subroutines.py

@@ -41,14 +41,14 @@ class Testconnectivity(unittest.TestCase):
     | (e) |
     0-----1
     """
-    # we test the connectivity matrix for nelx=2, nely=1, as above
-    # connectivity returns an array which relate the local
-    # numbering to global numbering. Thenumber of rows of connectivity
-    # is 4 since we use 4-nodes element and number of columns of connectivity
-    # is equal to number of element we have.
-    # for example compare the node numbering of element (e) with
+    # We test the connectivity matrix for nelx=2, nely=1, as above.
+    # Connectivity returns an array which relates the local
+    # numbering to the global numbering. The number of rows of the connectivity
+    # is 4 since we use 4-node elements and the number of the columns of
+    # the connectivity is equal to number of elements we have.
+    # For example, compare the node numbering of element (e) with
     # element (1): 0-->1, 1-->2, 2-->5, 3-->4
-    # so second column of connectivity will be 1,2,5,4
+    # So the second column of the connectivity will be 1,2,5,4
     def test_connectivity_2elements(self):
         IEN_test = [[0, 1], [1, 2], [4, 5], [3, 4]]
         IEN = FE.connectivity(2, 1)
@@ -77,7 +77,7 @@ def test_Dirichlet_BCs_2elements(self):
 
         # test flags
         for i in range(0, 4):
-            self.assertEqual(flags[i][0], 2)
+            self.assertEqual(flags[i][0], 1)
 
 
 class Testsetup_ID_LM(unittest.TestCase):