%auto
# Notebook setup
import numpy as np  # numerical package; subsequently typing "np." is translated into "numpy." and gives access to numpy commands. Type "numpy." into a code field and press <TAB> to see all possible commands in numpy.
import pylab   # Plotting package for more sophisticated plots
import csv     # Package for importing csv data files

# The below code loads units names into the name space for ease of use. Don't worry about it.
for s in units.amount_of_substance.trait_names():
    globals()[s] = getattr(units.amount_of_substance,s) 
for s in units.energy.trait_names():
    globals()[s] = getattr(units.energy,s)
for s in units.length.trait_names():
    globals()[s] = getattr(units.length,s)    
for s in units.mass.trait_names():
    globals()[s] = getattr(units.mass,s)
for s in units.pressure.trait_names():
    globals()[s] = getattr(units.pressure,s)
for s in units.temperature.trait_names():
    globals()[s] = getattr(units.temperature,s)
for s in units.time.trait_names():
    globals()[s] = getattr(units.time,s)
for s in units.power.trait_names():
    globals()[s] = getattr(units.power,s)
for s in units.pressure.trait_names():
    globals()[s] = getattr(units.pressure,s)
    
# Below, I create some empty dictionaries that will be populated later with variable descriptions, units and values of constants used in this worksheet.
udict = {}    # units
cdict = {}    # values of constants
docdict = {}    # variable descriptions

# Below, I define two functions that execute some code on the input given. The description of what it does is between the ''' marks. You can also view the description by typing "var2?" and Shift+Enter.

def var2(name,doc='',units=1,domain1='real',latexname='',value = 1234):
    '''
    Creates a symbolic variable in the given domain (standard:'real') and with the given 
    latexname. Further, it adds the string doc to docdict, the units to udict and the value
    to cdict. All entries are optional except for name.
    Usage example: 
        sage: var2('name','doc',watt/meter^2,'positive','N_{ame}',1.0)
    '''
    if len(latexname)>0:
        z = var(name,domain = domain1,latex_name = latexname)   
    else:
        z = var(name,domain = domain1) 
    if len(doc)>0:
        docdict[var(name)] = doc
    udict[var(name)] = var(name)*units
    if value != 1234:
        cdict[var(name)] = value
    return z
    
def units_check(eq):
    '''
    Checks whether all arguments are keys in udict and returns simplified units
    '''
    for blah in eq.arguments():
        udict[blah]
    return (eq.subs(udict)/eq).simplify_full()

# From http://www.engineeringtoolbox.com/weirs-flow-rate-d_592.html
q = var2('q', 'Flow rate', meter^3/second)
h = var2('h', 'Head on the weir', meter)
b = var2('b', 'Width of the weir', meter)
g = var2('g', 'Gravitational acceleration', meter^2/second, value = 9.81)
c_d = var2('c_d', 'Discharge constatn for the weir (calibrated)')
theta = var2('theta', 'V-notch angle, radians')

# Triangular v-notch
eq_q_vn = q == 8/15*c_d*sqrt(2*g)*tan(theta/2)*h^(5/2)
show(eq_q_vn)

vdict = cdict.copy()
vdict[theta] = 45/180*pi
vdict[c_d] = 1
P = plot(eq_q_vn.rhs().subs(vdict), (h, 0,1), frame = True, axes = False)
P.axes_labels(['h (m)', 'q (m$^3$ s$^{-1}$)'])
P.fontsize(16)
P