addNewDim.ncl

; addNewDim.ncl
;
; Provides functions for expanding and filling arrays of arbitrary rank. Combined the functions provide similar functionality to
;	array_append_record and conform_dims, but use ndtooned to avoid repeated variable creation/deletion needed for iterative 
;	array_append_record calls. The two functions, addNewDim and fillNewDim, are generally used together to combine data iteratively.
;	See example at end of this block for more information.
;
; Dependencies:
;	none
;	
; Included:
; 
; function addNewDim(variable, newDim, newDimName, newDimRank)
;	Use:
;		Adds an additional named dimension with coordinate data into an array of arbitrary size, retaining existing metadata.
;	Inputs:
;		variable 		,           ; variable of any dimensionality to be expanded
;		newDim[*]		,           ; either a scalar indicating the new dimension size, or an array with the coordinate data of the new dimension
;		newDimName[1]	: string    ; name for the new dimension
;		newDimRank[1]	: integer	; rank of the new dimension
;	Output:
;		An empty array of rank one higher than variable. The new dimension is added in the position indicated by newDimRank. This should generally be
;		0 (the slowest-varying dimension). The new dimension is named newDimName, has a size equal to the size of newDim, and the contents of newDim
;		are assigned as the coordinate variable for the new dimension. Dimensions other than the new dimension are identical to variable and retain 
;		all dimension names and coordinate information. All attributes of variable are copied to the new array as well.
;		
; function fillNewDim(variable, fill, index, newDimRank)
;	Use:
;		Adds data into an array of arbitrary rank at the specified index of a specified dimension.
;	Inputs:
;		variable 		: numeric, 	; variable we want to fill data into (generated by addNewDim)
;		fill 			: numeric, 	; data to be added into variable - dimensions must match those of variable other than the dim specified by fillRank
;		fillIndex[1]	: integer, 	; array index we want to fill
;		fillRank[1]		: integer 	; rank of the fill dimension
;	Output:
;		An array with the same data and metadata of variable, except data of the index at the left-most dimension replaced by fill.
;
; Example application:
;	Suppose we want to combine data from three different times of 0, 3 and 6 hours. The size of the data will be the same at each time, but is 
;	determined by inputs at run time [e.g., selecting surface temperature (2-D) or pressure (3-D) from an NWP model].
;	
;	*** Example code ***
;	
;	times = (/0,3,6/)
;	times@units = "hour"
;	n = dimsizes(times)
;	do t=0, n-1;	 
;		dataThisTime = //function to read data from file//
;	
;		; first time through, we create an array that will hold ALL of the data, based on the dimensions of the first piece of data
;		; the new dimensions is automatically named and assigned the appropriate coordinate data
;		if(t.eq.0) then
;			allData = addNewDim(dataThisTime, times, "time", 0)
;		end if
;	
;		; since allData isn't changing size each iteration, we can keep overwriting as new data is added
;		allData = fillNewDim(allData, dataThisTime, t, 0)
;	end do
;	
;	*** End example code ***
;	
;	After execution, allData will hold all of the data read in. Note that we did not need to know the dimensionality of the data being read in. 
;	If the original data is two-dimensional, the result will be three-dimensional. If the original data is three-dimensional, the result will be
;	four-dimensional, etc.
; 
;
; Created by:   Walter C. Kolczynski (Walter.Kolczynski@gmail.com)
; -- Last modified 18 July 2013
; 
; This work is released under the Creative Commons Attribution-ShareAlike 3.0 Unported License. Under this license, use and
;	modification are permitted as long as the authors are attributed and the distribution of any derived works is covered by
;	the same or similar license. For more details and full license text, see http://creativecommons.org/licenses/by-sa/3.0/
;	
; Use at your own risk.

load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/contrib/cd_string.ncl"

;
; Adds an additional named dimension with coordinate data into an array of arbitrary size, retaining existing metadata.
;
undef("addNewDim")
function addNewDim(                   	\\
			variable 		,           \\  ; variable to be expanded
			newDim[*]		,           \\  ; an array with the coordinate data of the new dimension
			newDimName[1]	: string,	\\  ; name for the new dimension
			newDimRank[1]	: integer	\\	; rank of the new dimension
			)
local oldDimSizes, nDims, err, oldDimRanks, oldDimNames, newDimSize, allDimSizes, newVariable, rank
begin

; get the existing dimensions
oldDimSizes = dimsizes(variable)
nDims = dimsizes(oldDimSizes)+1

if( newDimRank.gt.nDims-1 ) then
	print("Fatal (addNewDim): The new dimension rank is more than one larger than the number of existing dimensions")
	status_exit(1)
end if

if( newDimRank.lt.0 ) then
	print("Fatal (addNewDim): The new dimension rank must be a positive integer")
	status_exit(1)
end if

; turn off warnings to supress annoying get1Dindex_Exclude messages
err = NhlGetErrorObjectId()
setvalues err
	"errLevel" : "Fatal"          ; only report Fatal errors
end setvalues

oldDimRanks = get1Dindex_Exclude( ispan( 0, nDims-1, 1), newDimRank )

; turn warnings back on
setvalues err
	"errLevel" : "Warning"
end setvalues

oldDimNames = getvardims(variable)

; calculate how many dimensions there will be after expansion
newDimSize = dimsizes(newDim)

; create a new array to hold the data
allDimSizes = new(nDims, integer)
allDimSizes(oldDimRanks) = oldDimSizes
allDimSizes(newDimRank) = newDimSize

; expand the current variable to an additional dimension of the appropriate size
newVariable = new( allDimSizes, typeof(variable), getVarFillValue(variable) )
copy_VarAtts( variable, newVariable )

; name the new dimension
newVariable!newDimRank = newDimName

; assign coordinate data it to the new dimension
newVariable&$newDimName$ = newDim

; name and assign the coordinate arrays to the old dimensions
do i=0, nDims-2
	rank = oldDimRanks(i)
	newVariable!rank = oldDimNames(i)
	newVariable&$oldDimNames(i)$ = variable&$oldDimNames(i)$
end do

return newVariable

end 	; addNewDim

;
; Adds data into an array of arbitrary rank at the specified index of a specified dimension.
;
undef("fillNewDim")
procedure fillNewDim( 						\\
			variable 		,               \\  ; variable we want to fill data into
			fill 			,               \\  ; data we want to fill into variable
			fillIndex[1]    : integer, 		\\  ; array index we want to fill into
			fillRank[1]		: integer		\\ 	; the rank we want to fill into
			)
local varDimSizes, fillDimSizes, nDims, err, otherDimRanks, nBlocks, blockGap, blockSize, variable1D, fill1D, blocks, blockNumbers, toIndexes, fromIndexes
begin
; startCpuTime = get_cpu_time()

; read dimension sizes of the variable and fill
varDimSizes = dimsizes(variable)
fillDimSizes = dimsizes(fill)
nDims = dimsizes(varDimSizes)

; turn off warnings to supress annoying get1Dindex_Exclude messages
err = NhlGetErrorObjectId()
setvalues err
	"errLevel" : "Fatal"          ; only report Fatal errors
end setvalues

otherDimRanks = get1Dindex_Exclude( ispan( 0, nDims-1, 1), fillRank )

; turn warnings back on
err = NhlGetErrorObjectId()
setvalues err
	"errLevel" : "Warning"          ; only report Fatal errors
end setvalues

; make sure the dimension sizes of the new data match where we want to fill
if( any( varDimSizes(otherDimRanks).ne.fillDimSizes ) ) then
	print("Fatal (fillNewDim): Variable and fill not compatible!")
	printVarSummary(variable)
	printVarSummary(fill)
	status_exit(1)
end if

if(fillRank.ne.0) then
	nBlocks = product( varDimSizes( ispan(0, fillRank-1, 1) ) )
else
	nBlocks = 1
end if
blockGap = product( varDimSizes(fillRank:) )

if(fillRank.ne.nDims-1) then
	blockSize = product( varDimSizes( ispan(fillRank+1, nDims-1, 1) ) )
else
	blockSize = 1
end if

; flatten data into one dimension
variable1D = ndtooned(variable)
fill1D = ndtooned(fill)

blocks = ispan(0, blockSize-1, 1)
blockNumbers = ispan(0, nBlocks-1, 1)

toIndexes 	= ndtooned( conform_dims( (/nBlocks, blockSize/), blockNumbers*blockGap, 0 ) + 	\\
						conform_dims( (/nBlocks, blockSize/), blocks, 1) + 					\\
						conform_dims( (/nBlocks, blockSize/), fillIndex * blockSize, -1 ) )

fromIndexes = ndtooned( conform_dims( (/nBlocks, blockSize/), blockNumbers*blockSize, 0 ) + \\
						conform_dims( (/nBlocks, blockSize/), blocks, 1) )
delete(blocks)
delete(blockNumbers)
delete(nBlocks)
delete(blockGap)
delete(blockSize)

variable1D(toIndexes) = fill1D(fromIndexes)
delete(fill1D)
delete(toIndexes)
delete(fromIndexes)

; reformat data back into the appropriate dimensions
variable = (/ onedtond(variable1D, varDimSizes) /)
delete(variable1D)

; endCpuTime = get_cpu_time()
; cpuTime = endCpuTime - startCpuTime
; cpuTime@units = "seconds"
; print( "Dimensions index filled in " + cd_string( cpuTime, "%H hr %M min %S sec" ) )

end

;
; Returns the designated left-most indexes out of a variable with an arbitrary number of dimensions
;
undef("pullRecords")
function pullRecords( 								\\
					variable 		, 				\\  ; 
					indexes[*] 		: integer, 		\\  ;
					recordRank[1]	: integer		\\ 	;
					)
local varDimSizes, newDimSize, nDims, dimNames, err, otherDimRanks, otherDimNames, newDimSizes, variable1D, nBlocks, blockSize, pullIndexes, blocks, \\
		blockNumbers, recordNumber, newVariable, dimName, oldDim, newDim, rank
begin

; wallStart = stringtointeger( systemfunc("date +%s") )

; read dimension sizes of the variable and fill
varDimSizes = dimsizes(variable)
newDimSize = dimsizes(indexes)
nDims = dimsizes(varDimSizes)
dimNames = getvardims(variable)

if(recordRank.gt.nDims .or. recordRank.lt.0) then 
	print("Fatal (pullRecords): rank must be a non-negative integer less than the number of dimensions!")
	status_exit(1)
end if

if( any( indexes.gt.varDimSizes(recordRank) .or. indexes.lt.0 ) ) then
	print("Fatal (pullRecords): one or more indexes fall outside the bounds of the specified dimension!")
	status_exit(1)
end if

; turn off warnings to supress annoying get1Dindex_Exclude messages
err = NhlGetErrorObjectId()
setvalues err
	"errLevel" : "Fatal"          ; only report Fatal errors
end setvalues

otherDimRanks = get1Dindex_Exclude( ispan( 0, nDims-1, 1), recordRank )

; turn warnings back on
err = NhlGetErrorObjectId()
setvalues err
	"errLevel" : "Warning"          ; only report Fatal errors
end setvalues

otherDimNames = dimNames( otherDimRanks )

; create a new array to hold the data
newDimSizes = new(nDims, integer)
newDimSizes(otherDimRanks) = varDimSizes(otherDimRanks)
newDimSizes(recordRank) = newDimSize

; flatten data into one dimension
variable1D = ndtooned(variable)

if(recordRank.ne.0) then
	nBlocks = product( varDimSizes( ispan(0, recordRank-1, 1) ) )
else
	nBlocks = 1
end if
blockGap = product( varDimSizes(recordRank:) )

if(recordRank.ne.nDims-1) then
	blockSize = product( varDimSizes( ispan(recordRank+1, nDims-1, 1) ) )
else
	blockSize = 1
end if

pullIndexes = new( blockSize * nBlocks * newDimSize, integer )

blocks = ispan(0, blockSize-1, 1)
blockNumbers = ispan(0, nBlocks-1, 1)
recordNumber = ispan(0, newDimSize-1, 1)

pullIndexes = ndtooned( conform_dims( (/newDimSize, nBlocks, blockSize/), recordNumber*blockSize, 0 ) 	+ 	\\
						conform_dims( (/newDimSize, nBlocks, blockSize/), blockNumbers*blockGap, 1 ) 	+ 	\\
						conform_dims( (/newDimSize, nBlocks, blockSize/), blocks, 2) )
delete(blocks)
delete(blockNumbers)
delete(recordNumber)

newVariable = onedtond( variable1D(pullIndexes), newDimSizes )
copy_VarAtts( variable, newVariable )
delete(variable1D)
delete(pullIndexes)

; name the new dimension
dimName 	= dimNames(recordRank)
oldDim 		= variable&$dimName$
newDim 		= oldDim(indexes)
newVariable!recordRank = variable!recordRank

; assign coordinate data it to the new dimension
newVariable&$dimName$ = newDim

; name and assign the coordinate arrays to the old dimensions
do i=0, nDims-2
	rank = otherDimRanks(i)
	newVariable!rank = otherDimNames(i)
	newVariable&$otherDimNames(i)$ = variable&$otherDimNames(i)$
end do

; wallEnd  = stringtointeger( systemfunc("date +%s") )
; wallTime = wallEnd - wallStart
; wallTime@units = "seconds"
; print( newDimSize + " records pulled in " + cd_string( wallTime, "%H hr %M min %S sec" ) )

return newVariable

end 	; pullRecords

undef("test_addNewDim")
procedure test_addNewDim()
local data0, data1, data2, data
begin

data0 = (/  (/  000,    001,    002 /), \\
			(/  010,    011,    012 /)  /)
data1 = (/  (/  100,    101,    102 /), \\
			(/  110,    111,    112 /)  /)
data2 = (/  (/  200,    201,    202 /), \\
			(/  210,    211,    212 /)  /)
			
data0!0 = "B"
data0!1 = "C"
data0&C = (/0,1,2/)
data0&B = (/0,10/)
data1!0 = "B"
data1!1 = "C"
data1&C = (/0,1,2/)
data1&B = (/0,10/)
data2!0 = "B"
data2!1 = "C"
data2&C = (/0,1,2/)
data2&B = (/0,10/)

data = addNewDim( data0, (/ 0, 100, 200 /), "A", 0 )
fillNewDim( data, data0, 0, 0 )
fillNewDim( data, data1, 1, 0 )
fillNewDim( data, data2, 2, 0 )
print(data)

dataB = pullRecords( data, (/1,2/), 0 )
print(dataB)
delete(dataB)
dataB = pullRecords( data, (/1/), 1 )
print(dataB)
delete(dataB)
dataB = pullRecords( data, (/1,2/), 2 )
print(dataB)
delete(dataB)

delete(data)

data = addNewDim( data0, (/ 0, 100, 200 /), "A", 1 )
fillNewDim( data, data0, 0, 1 )
fillNewDim( data, data1, 1, 1 )
fillNewDim( data, data2, 2, 1 )
print(data)
delete(data)

data = addNewDim( data0, (/ 0, 100, 200 /), "A", 2 )
fillNewDim( data, data0, 0, 2 )
fillNewDim( data, data1, 1, 2 )
fillNewDim( data, data2, 2, 2 )
print(data)
delete(data)


end