def getNodeSourceQuality():
    epanet_proj = en.createproject()
    en.open(ph=epanet_proj,
            inpFile=example_1_path,
            rptFile='report.rpt',
            outFile='output.out')
    en.setqualtype(ph=epanet_proj,
                   qualType=1,
                   chemName='Chlorine',
                   chemUnits='mg/L',
                   traceNode=None)
    num_nodes = en.getcount(ph=epanet_proj, object=en.NODECOUNT)
    source_qual_list = []
    en.openH(ph=epanet_proj)
    en.initH(ph=epanet_proj, initFlag=0)
    en.openQ(ph=epanet_proj)
    en.initQ(ph=epanet_proj, saveFlag=1)
    #sets source_ind as booster, need to find a way to acknowledge pre-set boosters
    #en.setnodevalue(ph=epanet_proj, index=source_ind, property=en.SOURCETYPE, value=1)
    for i in range(1, num_nodes + 1):
        node_type = en.getnodevalue(ph=epanet_proj,
                                    index=i,
                                    property=en.SOURCETYPE)
        if node_type == 1:
            source_qual = en.getnodevalue(ph=epanet_proj,
                                          index=source_ind,
                                          property=en.SOURCEQUAL)
            print('Node (index %d) source quality: %5.2f' % (i, source_qual))
            source_qual_list.append(source_qual)
    print('End of getNodeSourceQuality function.')
    en.closeH(ph=epanet_proj)
    en.closeQ(ph=epanet_proj)
    en.close(ph=epanet_proj)
    clean_dir()
def test_water_age_sim():
    epanet_proj = en.createproject()
    en.open(ph=epanet_proj, inpFile=example_1_path, rptFile='report.rpt', outFile='output.out')
    en.setqualtype(ph=epanet_proj, qualType=2, chemName='Age', chemUnits='hours', traceNode=None)
    en.solveH(ph=epanet_proj)
    en.openQ(ph=epanet_proj)
    en.initQ(ph=epanet_proj, saveFlag=1)
    num_nodes = en.getcount(ph=epanet_proj, object=en.NODECOUNT)
    age_list = []
    while True:
        node_age_list = []
        en.runQ(ph=epanet_proj)
        t = en.nextQ(ph=epanet_proj)
        for i in range(1, num_nodes*1):
            node_qual = en.getnodevalue(ph=epanet_proj, index=i, property=en.QUALITY)
            node_age_list.append(node_qual)
        age_list.append(node_age_list)
        if t <= 0:
            break
    en.closeQ(ph=epanet_proj)
    en.closeH(ph=epanet_proj)
    en.close(ph=epanet_proj)
    assert age_list[26] == [1.0, 2.2141675704376946, 12.939125434025273, 24.44152992466322, 13.174235412569542,
                            24.441519659540887, 15.679376648181817, 21.97064181429266, 19.048343501261524, 1.0]
    clean_dir()
def test_set_qualtype():
    epanet_proj = en.createproject()
    en.open(ph=epanet_proj, inpFile=example_1_path, rptFile='report.rpt', outFile='output.out')
    en.setqualtype(ph=epanet_proj, qualType=2, chemName='Age', chemUnits='seconds', traceNode=None)
    qual_type = en.getqualtype(ph=epanet_proj)
    en.close(ph=epanet_proj)
    assert qual_type == [2, 0]
    clean_dir()
def test_water_quality():
    epanet_proj = en.createproject()
    en.open(ph=epanet_proj, inpFile=example_1_path, rptFile='report.rpt', outFile='output.out')
    en.setqualtype(ph=epanet_proj, qualType=1, chemName='Chlorine', chemUnits='mg/L', traceNode=None)
    num_nodes = en.getcount(ph=epanet_proj, object=en.NODECOUNT)
    num_links = en.getcount(ph=epanet_proj, object=en.LINKCOUNT)
    tlist = []
    node_cl_list = []
    link_cl_list = []
    en.openH(ph=epanet_proj)
    en.initH(ph=epanet_proj, initFlag=0)
    print('Printing hydraulic time step:')
    while True:
        en.runH(ph=epanet_proj)
        t = en.nextH(ph=epanet_proj)
        print(t)
        tlist.append(t)
        if t <= 0:
            break
    assert tlist == timesteps
    en.openQ(ph=epanet_proj)
    en.initQ(ph=epanet_proj, saveFlag=1)
    print('Printing chlorine concentration in nodes:')
    while True:
        en.runQ(ph=epanet_proj)
        t = en.nextQ(ph=epanet_proj)
        for i in range(1, num_nodes+1):
            node_qual = en.getnodevalue(ph=epanet_proj, index=i, property=en.QUALITY)
            print('Node %d: %5.2f' % (i, node_qual))
            node_cl_list.append(node_qual)
        if t <= 0:
            break
    print('Printing chlorine concentration in links:')
    while True:
        en.runQ(ph=epanet_proj)
        t = en.nextQ(ph=epanet_proj)
        for i in range(1, num_links+1):
            link_qual = en.getlinkvalue(ph=epanet_proj, index=i, property=en.QUALITY)
            print('Node %d: %5.2f' % (i, link_qual))
            link_cl_list.append(link_qual)
        if t <= 0:
            break
    en.closeH(ph=epanet_proj)
    en.closeQ(ph=epanet_proj)
    en.close(ph=epanet_proj)
    clean_dir()
def inject_chlorine(str_nodeID, booster_val):
    epanet_proj = en.createproject()
    en.open(ph=epanet_proj,
            inpFile=example_1_path,
            rptFile='report.rpt',
            outFile='output.out')
    en.setqualtype(ph=epanet_proj,
                   qualType=1,
                   chemName='Chlorine',
                   chemUnits='mg/L',
                   traceNode=None)
    num_nodes = en.getcount(ph=epanet_proj, object=en.NODECOUNT)
    num_links = en.getcount(ph=epanet_proj, object=en.LINKCOUNT)
    tlist = []
    en.openH(ph=epanet_proj)
    en.initH(ph=epanet_proj, initFlag=0)
    en.openQ(ph=epanet_proj)
    booster_node = en.getnodeindex(ph=epanet_proj, id=str_nodeID)
    print('Booster node index is %d' % (booster_node))
    #sets node as Mass Booster
    en.setnodevalue(ph=epanet_proj,
                    index=booster_node,
                    property=en.SOURCETYPE,
                    value=1)
    en.setnodevalue(ph=epanet_proj,
                    index=booster_node,
                    property=en.SOURCEQUAL,
                    value=booster_val)
    nodetype = en.getnodevalue(ph=epanet_proj,
                               index=booster_node,
                               property=en.SOURCETYPE)
    obj = QualityControl(1)  #seems to set nodetype to 1
    #obj.set_nodetype(1) runs automatically
    print(obj.get_nodetype())  #checking
    en.initQ(ph=epanet_proj, saveFlag=1)
    node_qual = en.getnodevalue(ph=epanet_proj,
                                index=booster_node,
                                property=en.SOURCEQUAL)
    print('Booster node source quality: %5.2f' % (node_qual))
    #print('End of inject chlorine function.')
    en.closeH(ph=epanet_proj)
    en.closeQ(ph=epanet_proj)
    en.close(ph=epanet_proj)
    clean_dir()
def setNodeSourceQuality(str_nodeID, booster_val):
    epanet_proj = en.createproject()
    en.open(ph=epanet_proj,
            inpFile=example_1_path,
            rptFile='report.rpt',
            outFile='output.out')
    en.setqualtype(ph=epanet_proj,
                   qualType=1,
                   chemName='Chlorine',
                   chemUnits='mg/L',
                   traceNode=None)
    num_nodes = en.getcount(ph=epanet_proj, object=en.NODECOUNT)
    num_links = en.getcount(ph=epanet_proj, object=en.LINKCOUNT)
    source_ind = en.getnodeindex(ph=epanet_proj, id=str_nodeID)
    '''en.openH(ph=epanet_proj)
    en.initH(ph=epanet_proj, initFlag=0)
    en.openQ(ph=epanet_proj)
    en.initQ(ph=epanet_proj, saveFlag=1)'''
    #setting source_ind as booster, need to find a way to acknowledge pre-set boosters
    #en.setnodevalue(ph=epanet_proj, index=source_ind, property=en.SOURCETYPE, value=1)
    node_type = en.getnodevalue(ph=epanet_proj,
                                index=source_ind,
                                property=en.SOURCETYPE)
    if node_type == 1:
        en.setnodevalue(ph=epanet_proj,
                        index=source_ind,
                        property=en.SOURCEQUAL,
                        value=booster_val)
        print('Node (index %d) source quality set to value %5.2f' %
              (source_ind, booster_val))
    else:
        print('Specified node is not designated as mass booster.')
    print('End of setNodeSourceQuality function.')
    '''en.closeH(ph=epanet_proj)
    en.closeQ(ph=epanet_proj)
    en.close(ph=epanet_proj)'''
    clean_dir()