updates to nodal power calc

import h5py

import numpy as np

import sys

def write_nodal_powers(file_name):

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

    """

    Write node-averaged power to an existing H5 file. A node is a quadrant of an assembly.

    Args:

        file_name (str):        Path to H5 file

    Assumes assemblies are symmetric across x and y dimensions (e.g., 17x17 assembly)

    """

    nStates = get_nStates(data)

    with h5py.File(file_name, 'r+') as data:

        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)]

            if 'pin_powers_nodal' in state.keys():

                del state['pin_powers_nodal']

        # 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 

            mult_mask = get_mult_mask(state)

            weighted_pin_powers = get_weighted_powers(inserted_pin_powers,mult_mask)

            node_mask = get_nodal_mask(state)

            nodal_power = get_nodal_powers(weighted_pin_powers,node_mask,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

            state['pin_powers_nodal'] = nodal_power

    data.close()

    return(nStates)

def get_nodal_mask_axial(core):

def get_assembly_dims(state):

    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)) ))

    pin_powers = np.squeeze(state['pin_powers'][:])

    n_asm_pin_x = pin_powers.shape[0]

    n_asm_pin_y = pin_powers.shape[1]

    n_asm_z = pin_powers.shape[2]

    return(mask)

    return(n_asm_pin_x,n_asm_pin_y,n_asm_z)

def get_assembly_midpoint(state):

    n_x,n_y,_ = get_assembly_dims(state)

    if n_x % 2 == 0:

        mid_x, mid_y = n_x/2, n_y/2

        odd = False

    else:

        mid_x, mid_y = (n_x-1)/2, (n_y-1)/2 

        odd = True

    return(mid_x,mid_y,odd)

def get_nodal_mask(core):

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

def get_nodal_mask(state):

    """

    Map of nodal regions for weighted pin power mapping

    """

    mid_x,mid_y,odd = get_assembly_midpoint(state)

    if odd:

        mask_dims = (mid_x+1,mid_y+1)

    else:

        mask_dims = (mid_x,mid_y)

    mask_0 = np.zeros(mask_dims, 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)

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))

    """

    Duplicate pin powers at the node boundaries

        via inline insertion of centerline row and column

    Passed with mult_mask to get weighted pin powers across the assembly

    """

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

    n_x,n_y,n_z = get_assembly_dims(state)

    mid_x,mid_y,odd = get_assembly_midpoint(state)

    pin_powers = np.squeeze(state['pin_powers'][:]) 

        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)

    if odd:

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

        inserted_pin_powers = np.zeros((n_x+1,n_y+1,n_z))

        for z in range(0,n_z):

            dummy = pin_powers[:,:,z]

            ins_row = dummy[mid_x,:]

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

            ins_col = dummy[:,mid_y]

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

            inserted_pin_powers[:,:,z] = dummy

    else: inserted_pin_powers = pin_powers

#    import pdb; pdb.set_trace()

    return(inserted_pin_powers)

def get_mult_mask(core):

def get_mult_mask(state):

    """

    Map of pin power weights across the inserted_pin_powers

        One assembly is broken into regions ((a,b),(c,d)) = (L,R)

    Weights at node boundaries are 0.5

    Weights at node vertex are 0.25

    """

    n_x,n_y,n_z  = get_assembly_dims(state)

    mid_x,mid_y,odd = get_assembly_midpoint(state)

    # 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

    # pin powers at node vertices weighted at 0.25

    if odd:

        mask_d = np.ones((mid_x+1,mid_y+1),dtype=float)

        mask_d[0,:],mask_d[:,0] = 0.5, 0.5

        mask_d[0,0] = 0.25

        mask_b = np.flipud(mask_d)

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

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

    else:

        mask = np.ones((n_x,n_y))

    return(mask)

def get_weighted_powers(pin_powers,mask):

    """

    Get weighted pin powers across the assembly

        pin_powers have been duplicated at node boundaries 

    """

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

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

    return(weighted_powers)

def get_nodal_powers(pin_powers,mask):

def get_nodal_powers(pin_powers,node_mask,mult_mask):

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

    """

    **** Apply mult_mask to inserted_pin_powers ****

    renormalize power

    """

    n_z = pin_powers.shape[2]

    nodal_power = np.zeros((2,2,n_z))

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

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

        power, weight, weighted_power = np.zeros((4,1)), np.zeros((4,1)), 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:

                    if node_mask[i,j] == n:

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

                        weight[n] += mult_mask[i,j] 

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

        weighted_power = np.divide(power,weight)

        nodal_power[:,:,z] = weighted_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)