Merge pull request #2565 from AllenInstitute/feature/2565-sf_new_average_mode

Add new mode for SF average operation

Author: t-b

Packages/MIES/MIES_Constants.ipf

@@ -2540,6 +2540,7 @@ StrConstant SF_OP_TPSS                = "tpss"
 StrConstant SF_OP_TPINST              = "tpinst"
 StrConstant SF_OP_TPBASE              = "tpbase"
 StrConstant SF_OP_TPFIT               = "tpfit"
+StrConstant SF_OP_EXTRACT             = "extract"
 ///@}
 
 StrConstant SF_PROPERTY_TABLE = "Table"

Packages/MIES/MIES_SweepFormula.ipf

@@ -98,7 +98,7 @@ Function/WAVE SF_GetNamedOperations()
 	                  SF_OP_MERGE, SF_OP_FIT, SF_OP_FITLINE, SF_OP_DATASET, SF_OP_SELECTVIS, SF_OP_SELECTCM, SF_OP_SELECTSTIMSET,       \
 	                  SF_OP_SELECTIVSCCSWEEPQC, SF_OP_SELECTIVSCCSETQC, SF_OP_SELECTRANGE, SF_OP_SELECTEXP, SF_OP_SELECTDEV,            \
 	                  SF_OP_SELECTEXPANDSCI, SF_OP_SELECTEXPANDRAC, SF_OP_SELECTSETCYCLECOUNT, SF_OP_SELECTSETSWEEPCOUNT,               \
-	                  SF_OP_SELECTSCIINDEX, SF_OP_SELECTRACINDEX, SF_OP_ANAFUNCPARAM, SF_OP_CONCAT, SF_OP_TABLE}
+	                  SF_OP_SELECTSCIINDEX, SF_OP_SELECTRACINDEX, SF_OP_ANAFUNCPARAM, SF_OP_CONCAT, SF_OP_TABLE, SF_OP_EXTRACT}
 
 	return wt
 End

Packages/MIES/MIES_SweepFormula_Executor.ipf

@@ -469,6 +469,9 @@ Function/WAVE SFE_FormulaExecutor(STRUCT SF_ExecutionData &exd, [variable srcLoc
 		case SF_OP_DATASET:
 			WAVE out = SFO_OperationDataset(exdop)
 			break
+		case SF_OP_EXTRACT:
+			WAVE out = SFO_OperationExtract(exdop)
+			break
 		case SF_OP_SELECTVIS:
 			WAVE out = SFOS_OperationSelectVis(exdop)
 			break

Packages/MIES/MIES_SweepFormula_Helpers.ipf

@@ -248,7 +248,9 @@ End
 /// @param copy              [optional, defaults to 0] If the returned data should be safe for modification (true) or is only read (false)
 /// @param[out] wvNote       [optional, defaults to None] Wave note of the dataset, useful for single result cases where you still need
 ///                          to query JSON wave note entries
-Function/WAVE SFH_GetArgumentAsWave(STRUCT SF_ExecutionData &exd, string opShort, variable argNum, [string defOp, WAVE/Z defWave, variable singleResult, variable expectedMinorType, variable expectedMajorType, variable copy, string &wvNote])
+/// @param resolveSelect     [optional, defaults to 0] If set then argument of the select type are automatically converted to sweep data, effectively the data() operation is applied
+///                          As with data, if the argument is an array of selects, each sweep result gets concatenated to a wref wave
+Function/WAVE SFH_GetArgumentAsWave(STRUCT SF_ExecutionData &exd, string opShort, variable argNum, [string defOp, WAVE/Z defWave, variable singleResult, variable expectedMinorType, variable expectedMajorType, variable copy, string &wvNote, variable resolveSelect])
 
 	variable checkExist, numArgs, checkMinorType, checkMajorType, result
 	string msg
@@ -273,12 +275,23 @@ Function/WAVE SFH_GetArgumentAsWave(STRUCT SF_ExecutionData &exd, string opShort
 	else
 		singleResult = !!singleResult
 	endif
-	copy = ParamIsDefault(copy) ? 0 : !!copy
+	copy          = ParamIsDefault(copy) ? 0 : !!copy
+	resolveSelect = ParamIsDefault(resolveSelect) ? 0 : !!resolveSelect
 
 	numArgs = SFH_GetNumberOfArguments(exd)
 
 	if(argNum < numArgs)
-		WAVE/WAVE input = SF_ResolveDatasetFromJSON(exd, argNum)
+		if(resolveSelect)
+			WAVE/Z/WAVE input = SFH_GetArgumentSelect(exd, argNum, doNotEnforce = 1)
+			if(WaveExists(input))
+				WAVE/WAVE data  = SFH_GetDataFromSelect(exd.graph, input)
+				WAVE      input = data
+			else
+				WAVE/WAVE input = SF_ResolveDatasetFromJSON(exd, argNum)
+			endif
+		else
+			WAVE/WAVE input = SF_ResolveDatasetFromJSON(exd, argNum)
+		endif
 
 		if(singleResult)
 			sprintf msg, "Argument #%d of operation %s: Too many input values", argNum, opShort
@@ -1067,16 +1080,24 @@ End
 ///        There is also a quick path for argNum >= numArgs, which is the case for e.g. data()
 ///        For that case numArgs is 0 and select is expected at argNum 0. Then the result of "select()" is
 ///        returned (as selectArray with a single element)
+///        If the doNotEnforce flag is set then a select type value at the arguments location is not enforced.
+///        If the argument is neither a single select or an array of selects then a null wave is returned.
 ///
 ///        selectArray is wave reference wave containing select composite wave reference waves with SELECTION, RANGE each.
 ///
 ///        This allows operations with selects as arguments to iterate over different selections given by the user
-Function/WAVE SFH_GetArgumentSelect(STRUCT SF_ExecutionData &exd, string opShort, variable argNum)
+Function/WAVE SFH_GetArgumentSelect(STRUCT SF_ExecutionData &exd, variable argNum, [variable doNotEnforce])
 
-	variable numArgs
-	string   type
+	variable numArgs, cond
+	string type
+
+	doNotEnforce = ParamIsDefault(doNotEnforce) ? 0 : !!doNotEnforce
 
 	numArgs = SFH_GetNumberOfArguments(exd)
+	if(argNum >= numArgs && doNotEnforce)
+		return $""
+	endif
+
 	if(argNum < numArgs)
 
 		WAVE/WAVE selectComp = SF_ResolveDatasetFromJSON(exd, argNum)
@@ -1087,15 +1108,31 @@ Function/WAVE SFH_GetArgumentSelect(STRUCT SF_ExecutionData &exd, string opShort
 			return selectArray
 		endif
 
-		SFH_ASSERT(DimSize(selectComp, ROWS) == 1, "Expected a single array")
+		cond = DimSize(selectComp, ROWS) == 1
+		if(doNotEnforce && !cond)
+			return $""
+		endif
+		SFH_ASSERT(cond, "Expected a single array")
 		WAVE array = selectComp[0]
-		SFH_ASSERT(IsTextWave(array), "Expected a text wave")
+		cond = IsTextWave(array)
+		if(doNotEnforce && !cond)
+			return $""
+		endif
+		SFH_ASSERT(cond, "Expected a text wave")
 
 		Make/FREE/WAVE/N=(DimSize(array, ROWS)) selectArray = SFH_AttemptDatasetResolve(WaveText(array, row = p), checkWithSFHAssert = 1)
 		for(WAVE/Z/WAVE selectComp : selectArray)
-			ASSERT(WaveExists(selectComp), "Expected select composite")
+			cond = WaveExists(selectComp)
+			if(doNotEnforce && !cond)
+				return $""
+			endif
+			ASSERT(cond, "Expected select composite")
 			type = JWN_GetStringFromWaveNote(selectComp, SF_META_DATATYPE)
-			SFH_ASSERT(!CmpStr(type, SF_DATATYPE_SELECTCOMP), "Expected select data as argument")
+			cond = !CmpStr(type, SF_DATATYPE_SELECTCOMP)
+			if(doNotEnforce && !cond)
+				return $""
+			endif
+			SFH_ASSERT(cond, "Expected select data as argument")
 		endfor
 
 		return selectArray
@@ -2070,3 +2107,53 @@ Function SFH_StoreAssertInfoExecutor(variable jsonId, variable srcLocId, string
 	info[%JSONPATH] = jsonPath
 	info[%STEP]     = num2istr(SF_STEP_EXECUTOR)
 End
+
+Function/WAVE SFH_GetDataFromSelect(string graph, WAVE/WAVE selectData)
+
+	WAVE/WAVE output = SFH_GetSweepsForFormula(graph, selectData, SF_OP_DATA)
+	if(!DimSize(output, ROWS))
+		DebugPrint("Call to SFH_GetSweepsForFormula returned no results")
+	endif
+
+	SFH_AddOpToOpStack(output, "", SF_OP_DATA)
+	SFH_ResetArgSetupStack(output, SF_OP_DATA)
+
+	return output
+End
+
+/// @brief Function returns a wave reference wave with one wave for each element of the dataset array
+///        Each wave contains the resolved dataset (as wave references)
+///
+/// @param exd               Execution data structure
+/// @param argNum            Argument index
+/// @param resolveSelect     [optional, defaults to 0] If set then argument of the select type are automatically converted to sweep data, effectively the data() operation is applied
+///                          As with data, if the argument is an array of selects, each sweep result gets concatenated to a wref wave
+/// @returns wref wave with one wave for each dataset array element. Each wave contains the resolved dataset as waveref wave.
+Function/WAVE SFH_GetDatasetArrayAsResolvedWaverefs(STRUCT SF_ExecutionData &exd, variable argNum, [variable resolveSelect])
+
+	variable numGroups
+
+	resolveSelect = ParamIsDefault(resolveSelect) ? 0 : !!resolveSelect
+
+	SFH_ASSERT(argNum < SFH_GetNumberOfArguments(exd), "operation has no argument number " + num2istr(argNum))
+	if(resolveSelect)
+		WAVE/Z/WAVE input = SFH_GetArgumentSelect(exd, argNum, doNotEnforce = 1)
+		if(WaveExists(input))
+			// In contrast to SFH_GetArgumentAsWave, resolve all selects from the array separately
+			numGroups = DimSize(input, ROWS)
+			SFH_ASSERT(numGroups >= 2, "First argument must be an array with at least two elements")
+			Make/FREE/WAVE/N=(numGroups) dataFromEachGroup = SFH_GetDataFromSelect(exd.graph, {input[p]})
+		endif
+	endif
+
+	if(!WaveExists(dataFromEachGroup))
+		WAVE/WAVE input = SF_ResolveDatasetFromJSON(exd, argNum)
+		SFH_ASSERT(IsTextWave(input[0]), "Expected a dataset")
+		WAVE/T elemRefs = input[0]
+		numGroups = DimSize(elemRefs, ROWS)
+		SFH_ASSERT(numGroups >= 2, "First argument must be an array with at least two elements")
+		Make/FREE/WAVE/N=(numGroups) dataFromEachGroup = SFH_AttemptDatasetResolve(elemRefs[p], checkWithSFHAssert = 1)
+	endif
+
+	return dataFromEachGroup
+End

Packages/MIES/MIES_SweepFormula_Operations.ipf

@@ -32,6 +32,7 @@ static StrConstant SF_OP_APFREQUENCY_X_TIME             = "time"
 
 static StrConstant SF_OP_AVG_INSWEEPS   = "in"
 static StrConstant SF_OP_AVG_OVERSWEEPS = "over"
+static StrConstant SF_OP_AVG_GROUPS     = "group"
 
 static StrConstant SF_OP_EPOCHS_TYPE_RANGE     = "range"
 static StrConstant SF_OP_EPOCHS_TYPE_NAME      = "name"
@@ -62,7 +63,7 @@ Function/WAVE SFO_OperationAnaFuncParam(STRUCT SF_ExecutionData &exd)
 	SFH_CheckArgumentCount(exd, SF_OP_ANAFUNCPARAM, 0, maxArgs = 2)
 
 	WAVE/T names      = SFH_GetArgumentAsWave(exd, SF_OP_ANAFUNCPARAM, 0, singleResult = 1)
-	WAVE/Z selectData = SFH_GetArgumentSelect(exd, SF_OP_DATA, 1)
+	WAVE/Z selectData = SFH_GetArgumentSelect(exd, 1)
 
 	WAVE/WAVE output = SFO_OperationAnaFuncParamIterate(exd.graph, names, selectData, SF_OP_ANAFUNCPARAM)
 
@@ -493,23 +494,63 @@ Function/WAVE SFO_OperationAvg(STRUCT SF_ExecutionData &exd)
 
 	numArgs = SFH_CheckArgumentCount(exd, opShort, 1, maxArgs = 2)
 
-	WAVE/WAVE input = SF_ResolveDatasetFromJSON(exd, 0)
-	mode = SFH_GetArgumentAsText(exd, opShort, 1, defValue = SF_OP_AVG_INSWEEPS, allowedValues = {SF_OP_AVG_INSWEEPS, SF_OP_AVG_OVERSWEEPS})
+	mode = SFH_GetArgumentAsText(exd, opShort, 1, defValue = SF_OP_AVG_INSWEEPS, allowedValues = {SF_OP_AVG_INSWEEPS, SF_OP_AVG_OVERSWEEPS, SF_OP_AVG_GROUPS})
+	if(!CmpStr(mode, SF_OP_AVG_INSWEEPS) || !CmpStr(mode, SF_OP_AVG_OVERSWEEPS))
+		WAVE/WAVE input = SFH_GetArgumentAsWave(exd, opShort, 0, resolveSelect = 1)
+		strswitch(mode)
+			case SF_OP_AVG_INSWEEPS:
+				WAVE/WAVE output = SFH_CreateSFRefWave(exd.graph, opShort, DimSize(input, ROWS))
+				output[] = SFO_OperationAvgImplIn(input[p])
+				SFH_TransferFormulaDataWaveNoteAndMeta(input, output, opShort, SF_DATATYPE_AVG)
+				return SFH_GetOutputForExecutor(output, exd.graph, opShort, clear = input)
 
-	strswitch(mode)
-		case SF_OP_AVG_INSWEEPS:
-			WAVE/WAVE output = SFH_CreateSFRefWave(exd.graph, opShort, DimSize(input, ROWS))
-			output[] = SFO_OperationAvgImplIn(input[p])
-			SFH_TransferFormulaDataWaveNoteAndMeta(input, output, opShort, SF_DATATYPE_AVG)
-			return SFH_GetOutputForExecutor(output, exd.graph, opShort, clear = input)
+			case SF_OP_AVG_OVERSWEEPS:
+				return SFO_OperationAvgImplOver(input, exd.graph, opShort)
 
-		case SF_OP_AVG_OVERSWEEPS:
-			return SFO_OperationAvgImplOver(input, exd.graph, opShort)
+			default:
+				FATAL_ERROR("Unhandled avg operation mode")
+		endswitch
+	elseif(!CmpStr(mode, SF_OP_AVG_GROUPS))
+		WAVE/WAVE dataFromEachGroup = SFH_GetDatasetArrayAsResolvedWaverefs(exd, 0, resolveSelect = 1)
+		WAVE/WAVE averagedGroup     = SFO_OperationAvgImplSweepGroups(dataFromEachGroup, exd.graph, opShort)
+		SFH_TransferFormulaDataWaveNoteAndMeta(dataFromEachGroup[0], averagedGroup, opShort, SF_DATATYPE_AVG)
 
-		default:
-			FATAL_ERROR("Unknown avg operation mode")
-	endswitch
+		return SFH_GetOutputForExecutor(averagedGroup, exd.graph, opShort)
+	else
+		FATAL_ERROR("Unhandled avg operation mode")
+	endif
+End
 
+static Function/WAVE SFO_OperationAvgImplSweepGroups(WAVE/WAVE sweepsFromEachSelection, string graph, string opShort)
+
+	variable numData, numMaxSweeps, numGroups, i, j
+	STRUCT RGBColor s
+
+	[s] = GetTraceColorForAverage()
+	Make/FREE/W/U traceColor = {s.red, s.green, s.blue}
+
+	numGroups = DimSize(sweepsFromEachSelection, ROWS)
+	Make/FREE/D/N=(numGroups) sweepCnts = DimSize(sweepsFromEachSelection[p], ROWS)
+	numMaxSweeps = WaveMax(sweepCnts)
+	WAVE/WAVE output = SFH_CreateSFRefWave(graph, opShort, numMaxSweeps)
+	for(i = 0; i < numMaxSweeps; i += 1)
+		Make/FREE/WAVE/N=(numGroups) avgSet
+		numData = 0
+		for(j = 0; j < numGroups; j += 1)
+			if(DimSize(sweepsFromEachSelection[j], ROWS) > i)
+				avgSet[numData] = WaveRef(sweepsFromEachSelection[j], row = i)
+				numData        += 1
+			endif
+		endfor
+		Redimension/N=(numData) avgSet
+		WAVE/WAVE avg = MIES_fWaveAverage(avgSet, 1, IGOR_TYPE_64BIT_FLOAT)
+		output[i] = avg[0]
+		JWN_SetWaveInWaveNote(output[i], SF_META_TRACECOLOR, traceColor)
+		JWN_SetNumberInWaveNote(output[i], SF_META_TRACETOFRONT, 1)
+		JWN_SetNumberInWaveNote(output[i], SF_META_LINESTYLE, 0)
+	endfor
+
+	return output
 End
 
 static Function/WAVE SFO_OperationAvgImplOver(WAVE/WAVE input, string graph, string opShort)
@@ -677,15 +718,9 @@ Function/WAVE SFO_OperationData(STRUCT SF_ExecutionData &exd)
 	variable i, numArgs
 
 	SFH_CheckArgumentCount(exd, SF_OP_DATA, 0, maxArgs = 1)
-	WAVE/WAVE selectData = SFH_GetArgumentSelect(exd, SF_OP_DATA, 0)
+	WAVE/WAVE selectData = SFH_GetArgumentSelect(exd, 0)
 
-	WAVE/WAVE output = SFH_GetSweepsForFormula(exd.graph, selectData, SF_OP_DATA)
-	if(!DimSize(output, ROWS))
-		DebugPrint("Call to SFH_GetSweepsForFormula returned no results")
-	endif
-
-	SFH_AddOpToOpStack(output, "", SF_OP_DATA)
-	SFH_ResetArgSetupStack(output, SF_OP_DATA)
+	WAVE output = SFH_GetDataFromSelect(exd.graph, selectData)
 
 	return SFH_GetOutputForExecutor(output, exd.graph, SF_OP_DATA)
 End
@@ -704,6 +739,24 @@ Function/WAVE SFO_OperationDataset(STRUCT SF_ExecutionData &exd)
 	return SFH_GetOutputForExecutor(output, exd.graph, SF_OP_DATASET)
 End
 
+// extract(<dataset>, index)
+Function/WAVE SFO_OperationExtract(STRUCT SF_ExecutionData &exd)
+
+	variable idx
+	string opShort = SF_OP_EXTRACT
+
+	SFH_CheckArgumentCount(exd, opShort, 2, maxArgs = 2)
+
+	WAVE/WAVE datasets = SFH_GetArgumentAsWave(exd, opShort, 0)
+	idx = SFH_GetArgumentAsNumeric(exd, opShort, 1)
+	SFH_ASSERT(idx >= 0 && idx < DimSize(datasets, ROWS), "index out of range")
+
+	WAVE/WAVE output = SFH_CreateSFRefWave(exd.graph, opShort, 1)
+	output[0] = datasets[idx]
+
+	return SFH_GetOutputForExecutor(output, exd.graph, opShort)
+End
+
 Function/WAVE SFO_OperationDerivative(STRUCT SF_ExecutionData &exd)
 
 	variable numArgs
@@ -808,7 +861,7 @@ Function/WAVE SFO_OperationEpochs(STRUCT SF_ExecutionData &exd)
 	endif
 
 	WAVE/Z/WAVE selectData      = $""
-	WAVE/Z/WAVE selectDataArray = SFH_GetArgumentSelect(exd, SF_OP_EPOCHS, 1)
+	WAVE/Z/WAVE selectDataArray = SFH_GetArgumentSelect(exd, 1)
 	if(WaveExists(selectDataArray))
 		SFH_ASSERT(DimSize(selectDataArray, ROWS) == 1, "Expected a single select specification")
 		WAVE/Z/WAVE selectDataComp = selectDataArray[0]
@@ -1090,7 +1143,7 @@ Function/WAVE SFO_OperationLabnotebook(STRUCT SF_ExecutionData &exd)
 	modeTxt = SFH_GetArgumentAsText(exd, SF_OP_LABNOTEBOOK, 2, allowedValues = allowedValuesMode, defValue = "DATA_ACQUISITION_MODE")
 	mode    = ParseLogbookMode(modeTxt)
 
-	WAVE/Z selectData = SFH_GetArgumentSelect(exd, SF_OP_LABNOTEBOOK, 1)
+	WAVE/Z selectData = SFH_GetArgumentSelect(exd, 1)
 
 	WAVE/T lbnKeys = SFH_GetArgumentAsWave(exd, SF_OP_LABNOTEBOOK, 0, expectedMajorType = IGOR_TYPE_TEXT_WAVE, singleResult = 1)
 

Packages/MIES/MIES_SweepFormula_Operations_TP.ipf

@@ -42,7 +42,7 @@ Function/WAVE SFOTP_OperationTP(STRUCT SF_ExecutionData &exd)
 		WAVE/Z ignoreTPs
 	endif
 
-	WAVE/Z selectData = SFH_GetArgumentSelect(exd, SF_OP_TP, 1)
+	WAVE/Z selectData = SFH_GetArgumentSelect(exd, 1)
 
 	WAVE/WAVE wMode = SF_ResolveDatasetFromJSON(exd, 0)
 	dataType = JWN_GetStringFromWaveNote(wMode, SF_META_DATATYPE)

Packages/MIES/MIES_SweepFormula_PSX.ipf

@@ -4760,7 +4760,7 @@ Function/WAVE PSX_OperationKernel(STRUCT SF_ExecutionData &exd)
 
 	SFH_CheckArgumentCount(exd, SF_OP_PSX_KERNEL, 0, maxArgs = 4)
 
-	WAVE/Z/WAVE selectDataCompArray = SFH_GetArgumentSelect(exd, SF_OP_PSX_KERNEL, 0)
+	WAVE/Z/WAVE selectDataCompArray = SFH_GetArgumentSelect(exd, 0)
 	SFH_ASSERT(WaveExists(selectDataCompArray), "Could not gather sweep data from select statement")
 
 	riseTau  = SFH_GetArgumentAsNumeric(exd, SF_OP_PSX_KERNEL, 1, defValue = 1, checkFunc = IsStrictlyPositiveAndFinite)
@@ -4910,7 +4910,7 @@ Function/WAVE PSX_OperationStats(STRUCT SF_ExecutionData &exd)
 
 	id = SFH_GetArgumentAsText(exd, SF_OP_PSX, 0, checkFunc = IsValidObjectName)
 
-	WAVE/Z/WAVE selectDataCompArray = SFH_GetArgumentSelect(exd, SF_OP_PSX_STATS, 1)
+	WAVE/Z/WAVE selectDataCompArray = SFH_GetArgumentSelect(exd, 1)
 	SFH_Assert(WaveExists(selectDataCompArray), "Missing select data")
 
 	WAVE allProps = PSX_GetAllStatsProperties()