updates to nodal power calc for pow3d

## ------------------------------------------------------ ##

# Author: Birdy

# Date: 1/10/20

# Description:

# Converts pin-by-pin powers into nodal powers

import h5py

import numpy as np

import os

class Reactor:

    def __init__(self,core,state):

        self.core = core

        self.state = state

        if 'pin_powers_nodal' in self.state.keys():

            del self.state['pin_powers_nodal']

        self.n_x, self.n_y, self.n_z, self.n_asm = self.state['pin_powers'][()].shape

        if self.n_x % 2 == 0:

            self.mid_x, self.mid_y = self.n_x/2, self.n_y/2

            self.odd = False

        else:

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

            self.odd = True

        self.pin_powers = state['pin_powers'][()]

        self.__get_inserted_pin_powers__()

        self.__get_mult_mask__()

        self.__get_nodal_mask__()

        self.__get_nodal_powers__()

        self.__perform_exit_checks__()

        self.state.create_dataset('pin_powers_nodal', data=self.nodal_powers)

    def __get_inserted_pin_powers__(self):

        """ Insert duplicated pin powers along node edges """

        if self.odd:

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

            for asm in range(0, self.n_asm):

                for z in range(0, self.n_z):

                    dummy = self.pin_powers[:,:,z,asm]

                    ins_row = dummy[self.mid_x,:]

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

                    ins_col = dummy[:,self.mid_y]

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

                    self.inserted_pin_powers[:,:,z,asm] = dummy

        else: self.inserted_pin_powers = self.pin_powers

    def __get_mult_mask__(self):

        """

        Define the pin power weights across the duplicated pin powers

        for a single axial slice

        Weights at node boundaries are 0.5 and at the vertex 0.25

        Regions are symmetric about x,y plane

        """

        if self.odd:

            mask_d = np.ones((self.mid_x+1,self.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)

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

            self.mult_mask = np.concatenate((L,R),axis=1)

        else: self.mult_mask = np.ones((self.n_x,self.n_y))

    def __get_nodal_mask__(self):

        """

        Define the nodal mask

        Axial slices are mapped to ((0,1),(2,3)) = (T,B)

        """

        if self.odd: mask_dims = (self.mid_x+1,self.mid_y+1)

        else: mask_dims = (self.mid_x,self.mid_y)

        mask_0 = np.zeros(mask_dims,dtype=int)

        mask_1 = np.ones(mask_dims,dtype=int)

        mask_2 = np.full(mask_dims,2,dtype=int)

        mask_3 = np.full(mask_dims,3,dtype=int)

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

    def __get_nodal_powers__(self):

        """

        Multiply element-by-element inserted pin powers; pin power weights

        Sum and group powers using nodal mask

        """

        self.nodal_powers = np.zeros( (self.core['core_map'].shape[0],self.core['core_map'].shape[1],4,self.n_z) )

        weighted_pin_powers = np.zeros(self.inserted_pin_powers.shape)

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

        self.core_map = self.core['core_map'][()]

        self.asm_map = self.core['core_map'][()]-1

        asm_i,asm_j = 0,0

        for core_asm in range(0,np.max(self.core_map)): # all assembly numbers

            for asm in range(0,len(np.argwhere(self.asm_map==core_asm))): # number of assembly matching type core_asm

                asm_i,asm_j = np.argwhere(self.asm_map==core_asm)[asm,0],np.argwhere(self.asm_map==core_asm)[asm,1] # coordinate of assembly matching type

                for z in range(0,self.n_z):

                    weighted_pin_powers = np.multiply(self.inserted_pin_powers[:,:,z,core_asm],self.mult_mask)

                    for n in range(0,4):

                        mask = self.node_mask == n

                        self.nodal_powers[asm_i,asm_j,n,z] = np.sum( np.multiply(weighted_pin_powers,mask) )

    def __perform_exit_checks__(self):

        """

        Check estimated nodal powers result in consistent:

        1. total power

           across entire core

        """

        total_power = self.core['rated_power'] * (self.state['power'][()] / 100.)

        core_avg_linear_heatrate = self.state['core_avg_linear_heatrate'][()]

        axial_mesh = self.core['axial_mesh'][:]

        mesh = axial_mesh[1:] - axial_mesh[:-1]

        ax_length = sum(mesh)

        local_power = np.zeros(self.nodal_powers.shape)

        for core_asm in range(np.max(self.core_map)):

            for z in range(0,self.n_z):

                for asm in range(0,len(np.argwhere(self.asm_map==core_asm))):

                    asm_i,asm_j = np.argwhere(self.asm_map==core_asm)[asm,0],np.argwhere(self.asm_map==core_asm)[asm,1]

                    local_power[asm_i,asm_j,:,z] = self.nodal_powers[asm_i,asm_j,:,z] * core_avg_linear_heatrate * mesh[z]

        est_total_power = 0.

        counter = 0

        for asm_i in range(0,self.core_map.shape[0]):

            for asm_j in range(0,self.core_map.shape[1]):

                est_total_power += np.sum( local_power[asm_i,asm_j,:,:] )

        est_total_power = est_total_power * 1e-6 # convert W to MW

        if not np.isclose(est_total_power, total_power, rtol=1e-3):

            print('rated power out of tolerance 1e-3')