import matplotlib.pyplot as plt
import numpy as np
%matplotlib inline



# Function repeated here else Notebook doesn't run
def pdrop(f=1,L=1,Q=1,D=1):
    delp = 8.0 * f * L * Q**2 / (9.81 * np.pi**2 * D**5)
    return delp

f = 0.02 * np.ones(2)
L = np.array([1000.0, 2000.0])
D = np.array([0.38, 0.3])
Q = np.array([300.0, 200.0]) / 3600.0 # To convert hours to seconds
p = np.zeros(3)
p[0] = 40 # This is point A
for i in range(2):
    p[i+1] = p[i] - pdrop( f = f[i], L = L[i], Q = Q[i], D = D[i] )
print p
plt.plot(p,'-bo')
plt.xticks(np.arange(3),['A','B','C'])
plt.xlabel('Node')
plt.ylabel('Pressure (m)')



def pdrop_triangle(delQ,Q0,Q1,Q2):
    f = 0.02 * np.ones(3)
    L = np.array([1600.0,1200.0,2000.0])
    D = 0.3*np.ones(3)
    Q = np.array([Q0+delQ,Q1+delQ,Q2-delQ])  # Discharges are adjusted according to delQ
    delp = np.zeros(3)
    for i in range(3):
        delp[i] = pdrop(f=f[i],L=L[i],D=D[i],Q=Q[i])
    pressure_drop = delp[0] + delp[1] - delp[2]
    return pressure_drop  # Pressure drop over the loop

from scipy.optimize import bisect
Q0,Q1,Q2 = 300.0/3600.0, 300.0/3600.0, 100.0/3600.0  # Q values need to be in m3/s
delQ = bisect( pdrop_triangle, a=Q2, b=-Q0, args=(Q0,Q1,Q2) )
print delQ

print 'pressure drop over triangle:',pdrop_triangle(delQ, Q0, Q1, Q2)
print 'The discharge delQ is',delQ*3600,'m3/h'

def pdrop_triangle(delQ,Q0,Q1,Q2,full_output=False):
    f = 0.02 * np.ones(3)
    L = np.array([1600.0,1200.0,2000.0])
    D = 0.3 * np.ones(3)
    Q = np.array([Q0+delQ,Q1+delQ,Q2-delQ])  # Discharges are adjusted according to delQ
    delp = np.zeros(3)
    for i in range(3):
        delp[i] = pdrop(f=f[i],L=L[i],D=D[i],Q=Q[i])
    pressure_drop = delp[0] + delp[1] - delp[2]
    if full_output:
        return pressure_drop, Q, delp
    return pressure_drop  # Pressure drop over the loop
delQ = bisect( pdrop_triangle, Q2, -Q0, args=(Q0,Q1,Q2) )
pd, Q, delp = pdrop_triangle(delQ,Q0,Q1,Q2,True)
print 'pressure drop over loop: ',pd
print 'discharges in pipes: ',Q * 3600 # To convert back to m3/hour
print 'pressure drops over pipes: ',delp
print 'pC through pipes 0-1: ',40.0 - delp[0] - delp[1]
print 'pC through pipe 2: ',40.0 - delp[2]









def pdrop(f=1,L=1,Q=1,D=1):
    delh = 8.0 * f * L * Q**2 / (9.81 * np.pi**2 * D**5)
    return delh
pdrop(f=0.02, L=1000, Q=200.0/60.0/60.0, D=0.3)

f = 0.02 * np.ones(10)
L = 1000 * np.ones(10)
D = np.ones(10)
D[:5] = 0.38
D[5:] = 0.3
Q = np.arange(500,0,-50) / 3600.0
p = np.zeros(11)
p[0] = 40 # This is point A
for i in range(10):
    p[i+1] = p[i] - pdrop( f = f[i], L = L[i], Q = Q[i], D = D[i] )
plt.plot(p,'-bo')
plt.xticks(np.arange(11),['A','B','C','D','E','F','G','H','I','J','K'])
plt.ylabel('Pressure (m)')
plt.xlabel('Node')

def pdrop_rect(delQ,Q0,Q1,Q2,Q3,full_output=False):
    f = 0.02 * np.ones(4)
    L = np.array([1600.0,1200.0,1600.0,1200.0])
    D = 0.3 * np.ones(4)
    Q = np.array([Q0+delQ,Q1+delQ,Q2-delQ,Q3-delQ])  # Discharges are adjusted according to delQ
    delp = np.zeros(4)
    for i in range(4):
        delp[i] = pdrop(f=f[i],L=L[i],D=D[i],Q=Q[i])
    pressure_drop = delp[0] + delp[1] - delp[2] - delp[3]
    if full_output:
        return pressure_drop, Q, delp
    return pressure_drop  # Pressure drop over the loop
Q0,Q1,Q2,Q3 = 300.0/3600, 100.0/3600, 200.0/3600, 200.0/3600
delQ = bisect( pdrop_rect, -Q0, Q3, args=(Q0,Q1,Q2,Q3) )
pd, Q, delp = pdrop_rect(delQ,Q0,Q1,Q2,Q3,True)
print 'pressure drop over loop: ',pd
print 'discharges in pipes: ',Q * 3600 # To convert back to m3/hour
print 'pressure drops over pipes: ',delp
print 'pressure in C through pipes 0-1: ',40.0 - delp[0] - delp[1]
print 'pressure in C through pipes 2-3: ',40.0 - delp[2] - delp[3]

def pdrop_loop1(delQ,Q0,Q1,Q4,full_output=False):
    f = 0.02 * np.ones(3)
    L = np.array([1600.0,1200.0,2000.0])
    D = 0.3 * np.ones(3)
    Q = np.array([Q0+delQ,Q1+delQ,Q4-delQ])  # Discharges are adjusted according to delQ
    delp = np.zeros(3)
    for i in range(3):
        delp[i] = pdrop(f=f[i],L=L[i],D=D[i],Q=Q[i])
    pressure_drop = delp[0] + delp[1] - delp[2]
    if full_output:
        return pressure_drop, Q, delp
    return pressure_drop  # Pressure drop over the loop
Q0,Q1,Q4 = 300.0 / 3600, 100.0 / 3600, 100.0 / 3600
delQ1 = bisect( pdrop_loop1, -Q0, Q4, args=(Q0,Q1,Q4) )
print 'delQ1: ',delQ1,delQ1*3600

def pdrop_loop2(delQ,Q2,Q3,Q4,full_output=False):
    f = 0.02 * np.ones(3)
    L = np.array([1200.0,1600.0,2000.0])
    D = 0.3 * np.ones(3)
    Q = np.array([Q2-delQ,Q3-delQ,Q4+delQ])  # Discharges are adjusted according to delQ
    delp = np.zeros(3)
    for i in range(3):
        delp[i] = pdrop(f=f[i],L=L[i],D=D[i],Q=Q[i])
    pressure_drop = -delp[0] - delp[1] + delp[2]
    if full_output:
        return pressure_drop, Q, delp
    return pressure_drop  # Pressure drop over the loop

Q2,Q3,Q4 = 100.0 / 3600, 100.0 / 3600, 100.0 / 3600
delQ2 = bisect( pdrop_loop2, -Q4, Q3, args=(Q2,Q3,Q4) )
print delQ2, delQ2*3600

Q0 = 300.0 / 3600
Q1 = 100.0 / 3600
Q2 = 100.0 / 3600
Q3 = 100.0 / 3600
Q4 = 100.0 / 3600

for i in range(10):
    delQ1 = bisect( pdrop_loop1, -Q0, Q4, args=(Q0,Q1,Q4) )
    Q0 = Q0 + delQ1
    Q1 = Q1 + delQ1
    Q4 = Q4 - delQ1
    delQ2 = bisect( pdrop_loop2, -Q4, Q3, args=(Q2,Q3,Q4) )
    Q2 = Q2 - delQ2
    Q3 = Q3 - delQ2
    Q4 = Q4 + delQ2
    print np.array([Q0,Q1,Q2,Q3,Q4])*3600
    
delQ1, Q, delp = pdrop_loop1(delQ1,Q0,Q1,Q4,True)
pC = 40 - delp[2]
print 'Discharges in the pipes Q0,Q1,Q2,Q3,Q4:',Q0*3600,Q1*3600,Q2*3600,Q3*3600,Q4*3600
print 'Pressure in point C: ',pC