Add nodal power calculator script

import h5py

import numpy as np

def write_nodal_powers(file_name):

    data = h5py.File(file_name, 'r+')

    nStates = get_nStates(data)

    for istate in range(nStates):

        # Point to core and state data

        core = data['CORE']

        state = data['STATE_{:04d}'.format(istate+1)]

        # get pin powers

        # insert duplicated pin powers along node edges

        inserted_pin_powers = insert_pin_powers(state)

        # get multiplier mask

        mult_mask = get_mult_mask(core)

        # get individual pin contributions with duplicates 

        weighted_pin_powers = get_weighted_powers(inserted_pin_powers,mult_mask)

        # sum and assign to nodes

        node_mask = get_nodal_mask(core)

        nodal_power = get_nodal_powers(weighted_pin_powers,node_mask)

        data['STATE_{:04d}'.format(istate+1)]['pin_powers_nodal'] = nodal_power

    data.close()

def get_nStates(data):

    nStates = 0

    for key in data.keys():

        if 'STATE' in key: nStates+=1

    return(nStates)

def get_nodal_mask_axial(core):

    mask_0 = np.zeros((9,9,49,1),dtype=int)

    mask_1 = np.ones((mask_0.shape),dtype=int)

    mask_2 = np.full((mask_0.shape),2,dtype=int)

    mask_3 = np.full((mask_0.shape),3,dtype=int)

    mask = np.vstack(( np.hstack((mask_0,mask_1)),np.hstack((mask_2,mask_3)) ))

    return(mask)

def get_nodal_mask(core):

    mask_0 = np.zeros((9,9,1),dtype=int)

    mask_1 = np.ones((mask_0.shape),dtype=int)

    mask_2 = np.full((mask_0.shape),2,dtype=int)

    mask_3 = np.full((mask_0.shape),3,dtype=int)

    mask = np.vstack(( np.hstack((mask_0,mask_1)),np.hstack((mask_2,mask_3)) ))

    return(mask)

def insert_pin_powers(state):

    # duplicate pin powers at node edges

    pin_powers = state['pin_powers'][:]

    inserted_pin_powers = np.zeros((18,18,49,1))

    for z in range(0,pin_powers.shape[2]):

        dummy = pin_powers[:,:,z,0]

        ins_row = dummy[8]

        dummy = np.insert(dummy,8,ins_row,axis=0)

        ins_col = dummy[:,8]

        dummy = np.insert(dummy,8,ins_col,axis=1)

        inserted_pin_powers[:,:,z,0] = dummy

#    import pdb; pdb.set_trace()

    return(inserted_pin_powers)

def get_mult_mask(core):

    # core = ((a,b),(c,d)) = ((L,R))

    # pin powers at node edges weighted at 0.5

    mask_d = np.ones((9,9,1),dtype=float)

    mask_d[0,:] = 0.5

    mask_d[:,0] = 0.5

    mask_b = np.flipud(mask_d)

    R = np.concatenate((mask_b,mask_d),axis=0)

    L = np.fliplr(R[:,:])

    mask = np.concatenate((L,R),axis=1)

    return(mask)

def get_weighted_powers(pin_powers,mask):

    weighted_powers = np.zeros((pin_powers.shape))

    for z in range(0,pin_powers.shape[2]):

        weighted_powers[:,:,z] = np.multiply(pin_powers[:,:,z],mask)

    return(weighted_powers)

def get_nodal_powers(pin_powers,mask):

    nodal_power = np.zeros((2,2,49,1))

    for z in range(0,pin_powers.shape[2]):

        power = np.zeros((4,1))

        for n in range(0,4):

            for i in range(pin_powers.shape[0]):

                for j in range(pin_powers.shape[1]):

                    if mask[i,j] == n:

                        power[n] += pin_powers[i,j,z]

        nodal_power[:,:,z,0] = power.reshape(2,2) 

    return(nodal_power)

def main():

    #file_name = '/lustre/hydra/cades-nsed-veracs/ubp/scratch/nodal_power_calc/Enrichment.h5'

    file_name = sys.argv[1]

    write_nodal_powers(file_name)

if __name__ == '__main__':

    main()