Added contour plot feature

ThermoPyle/ThermoPyle.pyx

@@ -9,13 +9,17 @@
 import json
 from copy import deepcopy
 from itertools import permutations
+from random import randint
 # from numba import jit, jitclass, int64, float64, void
 
 import arrow
 import CoolProp.CoolProp as CP
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import matplotlib
+import matplotlib.cm as cm
+import matplotlib.mlab as mlab
 import pyrsistent as pyr
 from cytoolz import get
 from mpl_toolkits.mplot3d import Axes3D
@@ -393,11 +397,11 @@ def fluid_plot(fluid: Union[CSVFluid, ThermoFluid]) -> None:
         edgecolors="none")
 
     # Set the Labels
-    ax.set_xlabel("{0} [{1}]".format(fluid.vars[0], fluid.units[0]))
-    ax.set_ylabel("{0} [{1}]".format(fluid.vars[1], fluid.units[1]))
-    ax.set_zlabel("{0} [{1}]".format(fluid.vars[2], fluid.units[2]))
-    ax.set_title("{0} and {1} vs {2} of {3}".format(*fluid.vars, fluid.fluid))
-    plt.show()
+    ax.set_xlabel("{0} [{1}]".format(fluid.xvar, fluid.units[fluid.xvar]))
+    ax.set_ylabel("{0} [{1}]".format(fluid.yvar, fluid.units[fluid.yvar]))
+    ax.set_zlabel("{0} [{1}]".format(fluid.zvar, fluid.units[fluid.zvar]))
+    ax.set_title("{0} and {1} vs {2} of {3}".format(fluid.xvar, fluid,yvar, fluid.zvar, fluid.fluid))
+    plt.show(fig)
 
 
 def rescale(oldrange: List[Union[float, int]],
@@ -412,3 +416,34 @@ def rescale(oldrange: List[Union[float, int]],
 
     """
     return lambda x: (newrange[1] - newrange[0]) / (oldrange[1] - oldrange[0]) * (x - oldrange[0]) + newrange[0]
+
+
+def fluid_contour_plot(fluid: Union[CSVFluid, ThermoFluid], contour: Text) -> None:
+    """
+    Does what it says on the tin. Makes a Contour Plot of the Fluid.
+
+    Paramaters:
+        fluid (Union[CSVFluid, ThermoFluid]): What fluid object to get data from.
+        contour (Text): What contour to plot. Must be a variable in fluid.vars
+    """
+    try:
+        assert contour in set(fluid.vars)
+    except AssertionError:
+        ValueError("contour needs to be a variable of fluid")
+
+    matplotlib.rcParams['xtick.direction'] = 'out'
+    matplotlib.rcParams['ytick.direction'] = 'out'
+    fig = plt.figure(randint(1,10**4))
+    fignum = fig.number
+    X = np.array(pd.Series(fluid.data[fluid.xvar]).unique())
+    Y = np.array(pd.Series(fluid.data[fluid.yvar]).unique())
+    Z = np.array(fluid.data.pivot(index=fluid.xvar, columns=fluid.yvar)[contour])
+    CS = plt.contour(X, Y.T, Z, cmap=fluid.colorMap)
+    fmt = {}
+    for x in CS.levels:
+        fmt[x] = "{:.4g}".format(x) + " " + fluid.units[contour]
+    plt.clabel(CS, fontsize=10, inline=1, fmt=fmt)
+    plt.xlabel("{0} [{1}]".format(fluid.xvar, fluid.units[fluid.xvar]))
+    plt.ylabel("{0} [{1}]".format(fluid.yvar, fluid.units[fluid.yvar]))
+    plt.title("{0} and {1} vs {2} of {3}".format(fluid.xvar, fluid.yvar, contour, fluid.fluid))
+    plot.show(fig)

docs/source/conf.py

@@ -49,7 +49,7 @@
 
 autodoc_mock_imports = [
     "pyrsistent", "arrow", "matplotlib", "mpl_toolkits", "matplotlib.pyplot",
-    "pandas", "cytoolz", "numpy", "mpl_toolkits.mplot3d", "Cython"
+    "pandas", "cytoolz", "numpy", "mpl_toolkits.mplot3d", "Cython", "matplotlib.mlab", "matplotlib.cm"
 ]
 
 # Add any paths that contain templates here, relative to this directory.