From 030ffb39531be507650d0013794085836571f2ca Mon Sep 17 00:00:00 2001
From: Jaehwan Ryu <jryu@uncharted.software>
Date: Thu, 2 Jan 2025 10:44:38 -0500
Subject: [PATCH] Calibrate ensemble tweaks (#5871)

Co-authored-by: Tom Szendrey <T.Szendrey@hotmail.com>
---
 .../widgets/tera-chart-settings.vue           | 46 +++++-----
 .../calibrate-ensemble-util.ts                | 26 +++++-
 ...ra-calibrate-ensemble-ciemss-drilldown.vue | 44 ++++++----
 .../tera-calibrate-ensemble-node-ciemss.vue   | 85 +++++++++++++++----
 4 files changed, 141 insertions(+), 60 deletions(-)

diff --git a/packages/client/hmi-client/src/components/widgets/tera-chart-settings.vue b/packages/client/hmi-client/src/components/widgets/tera-chart-settings.vue
index 1dd138f4a25..a60ede2df49 100644
--- a/packages/client/hmi-client/src/components/widgets/tera-chart-settings.vue
+++ b/packages/client/hmi-client/src/components/widgets/tera-chart-settings.vue
@@ -1,6 +1,29 @@
 <template>
 	<div class="chart-settings">
 		<h5>{{ title }}</h5>
+		<template v-if="type === ChartSettingType.VARIABLE_ENSEMBLE">
+			<tera-checkbox
+				:disabled="selectedOptions.length === 0"
+				label="Show individual models"
+				:model-value="Boolean(ensembleChartOptions?.showIndividualModels)"
+				@update:model-value="toggleEnsembleChartOption('showIndividualModels', $event)"
+			/>
+			<!-- Disabling following two checkboxes for now since their functionalities aren't implemented yet -->
+			<!-- <tera-checkbox
+				class="pl-5"
+				:disabled="selectedOptions.length === 0"
+				label="Relative to ensemble"
+				:model-value="Boolean(ensembleChartOptions.relativeToEnsemble)"
+				@update:model-value="toggleEnsembleChartOption('relativeToEnsemble', $event)"
+			/>
+			<tera-checkbox
+				v-if="isSimulateEnsembleSettings"
+				label="Show individual models with weights"
+				:disabled="selectedOptions.length === 0"
+				:model-value="Boolean(ensembleChartOptions.showIndividualModelsWithWeight)"
+				@update:model-value="toggleEnsembleChartOption('showIndividualModelsWithWeight', $event)"
+			/> -->
+		</template>
 		<tera-chart-control
 			:chart-config="{
 				selectedRun: 'fixme',
@@ -56,29 +79,6 @@
 			@open="$emit('open', s)"
 			@remove="$emit('remove', s.id)"
 		/>
-		<template v-if="type === ChartSettingType.VARIABLE_ENSEMBLE">
-			<tera-checkbox
-				:disabled="selectedOptions.length === 0"
-				label="Show individual models"
-				:model-value="Boolean(ensembleChartOptions?.showIndividualModels)"
-				@update:model-value="toggleEnsembleChartOption('showIndividualModels', $event)"
-			/>
-			<!-- Disabling following two checkboxes for now since their functionalities aren't implemented yet -->
-			<!-- <tera-checkbox
-				class="pl-5"
-				:disabled="selectedOptions.length === 0"
-				label="Relative to ensemble"
-				:model-value="Boolean(ensembleChartOptions.relativeToEnsemble)"
-				@update:model-value="toggleEnsembleChartOption('relativeToEnsemble', $event)"
-			/>
-			<tera-checkbox
-				v-if="isSimulateEnsembleSettings"
-				label="Show individual models with weights"
-				:disabled="selectedOptions.length === 0"
-				:model-value="Boolean(ensembleChartOptions.showIndividualModelsWithWeight)"
-				@update:model-value="toggleEnsembleChartOption('showIndividualModelsWithWeight', $event)"
-			/> -->
-		</template>
 	</div>
 </template>
 
diff --git a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/calibrate-ensemble-util.ts b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/calibrate-ensemble-util.ts
index 3d2bb2cfd46..c87750757f4 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/calibrate-ensemble-util.ts
+++ b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/calibrate-ensemble-util.ts
@@ -8,10 +8,10 @@ import {
 	getSimulation,
 	parseEnsemblePyciemssMap
 } from '@/services/models/simulation-service';
-import { EnsembleModelConfigs } from '@/types/Types';
+import { EnsembleModelConfigs, ModelConfiguration } from '@/types/Types';
 import { WorkflowNode } from '@/types/workflow';
 import { getActiveOutput } from '@/components/workflow/util';
-import { CalibrateMap } from '@/services/calibrate-workflow';
+import { CalibrateMap, setupModelInput } from '@/services/calibrate-workflow';
 import {
 	CalibrateEnsembleCiemssOperationState,
 	CalibrateEnsembleMappingRow,
@@ -148,6 +148,28 @@ export async function fetchOutputData(preForecastId: string, postForecastId: str
 	};
 }
 
+export async function fetchModelConfigurations(
+	nodeInputs: WorkflowNode<CalibrateEnsembleCiemssOperationState>['inputs']
+) {
+	const allModelOptions: any[][] = [];
+	const allModelConfigurations: ModelConfiguration[] = [];
+	const modelConfigurationIds: string[] = [];
+	nodeInputs.forEach((ele) => {
+		if (ele.value && ele.type === 'modelConfigId') modelConfigurationIds.push(ele.value[0]);
+	});
+	if (!modelConfigurationIds) return null;
+
+	// Model configuration input
+	await Promise.all(
+		modelConfigurationIds.map(async (id) => {
+			const { modelConfiguration, modelOptions } = await setupModelInput(id);
+			if (modelConfiguration) allModelConfigurations.push(modelConfiguration);
+			if (modelOptions) allModelOptions.push(modelOptions);
+		})
+	);
+	return { allModelConfigurations, allModelOptions };
+}
+
 // Build chart data by adding variable translation map to the given output data
 export function buildChartData(
 	outputData: {
diff --git a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-ciemss-drilldown.vue b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-ciemss-drilldown.vue
index d089e20d3a9..1fcd393a063 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-ciemss-drilldown.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-ciemss-drilldown.vue
@@ -345,7 +345,7 @@ import TeraInputNumber from '@/components/widgets/tera-input-number.vue';
 import AccordionTab from 'primevue/accordiontab';
 import Accordion from 'primevue/accordion';
 import Dropdown from 'primevue/dropdown';
-import { setupDatasetInput, setupCsvAsset, setupModelInput, parseCsvAsset } from '@/services/calibrate-workflow';
+import { setupDatasetInput, setupCsvAsset, parseCsvAsset } from '@/services/calibrate-workflow';
 import TeraDrilldown from '@/components/drilldown/tera-drilldown.vue';
 import TeraDrilldownSection from '@/components/drilldown/tera-drilldown-section.vue';
 import TeraSaveDatasetFromSimulation from '@/components/dataset/tera-save-dataset-from-simulation.vue';
@@ -375,7 +375,7 @@ import VegaChart from '@/components/widgets/VegaChart.vue';
 import { ChartSettingType, CiemssPresetTypes, DrilldownTabs } from '@/types/common';
 import { useCharts } from '@/composables/useCharts';
 import { useChartSettings } from '@/composables/useChartSettings';
-import { deleteAnnotation } from '@/services/chart-settings';
+import { deleteAnnotation, updateChartSettingsBySelectedVariables } from '@/services/chart-settings';
 import { DataArray } from '@/utils/stats';
 import { GroupedDataArray } from '@/services/charts';
 import {
@@ -393,7 +393,8 @@ import {
 	fetchOutputData,
 	buildChartData,
 	getEnsembleErrorData,
-	EnsembleErrorData
+	EnsembleErrorData,
+	fetchModelConfigurations
 } from './calibrate-ensemble-util';
 
 const props = defineProps<{
@@ -550,6 +551,22 @@ const setPresetValues = (data: CiemssPresetTypes) => {
 	}
 };
 
+const initDefaultChartSettings = (state: CalibrateEnsembleCiemssOperationState) => {
+	const mappedEnsembleVariables = knobs.value.ensembleMapping.map((c) => c.newName);
+	if (_.isEmpty(state.chartSettings)) {
+		state.chartSettings = updateChartSettingsBySelectedVariables(
+			state.chartSettings ?? [],
+			ChartSettingType.VARIABLE_ENSEMBLE,
+			mappedEnsembleVariables
+		);
+		state.chartSettings = updateChartSettingsBySelectedVariables(
+			state.chartSettings,
+			ChartSettingType.ERROR_DISTRIBUTION,
+			mappedEnsembleVariables
+		);
+	}
+};
+
 const runEnsemble = async () => {
 	if (!datasetId.value || !currentDatasetFileName.value) return;
 
@@ -590,26 +607,17 @@ const runEnsemble = async () => {
 		state.currentProgress = 0;
 		state.inProgressCalibrationId = response?.simulationId;
 		state.inProgressForecastId = '';
+		// Add default chart settings based on the ensemble mapping on the first run
+		initDefaultChartSettings(state);
 		emit('update-state', state);
 	}
 };
 
 onMounted(async () => {
-	allModelConfigurations.value = [];
-	const modelConfigurationIds: string[] = [];
-	props.node.inputs.forEach((ele) => {
-		if (ele.value && ele.type === 'modelConfigId') modelConfigurationIds.push(ele.value[0]);
-	});
-	if (!modelConfigurationIds) return;
-
-	// Model configuration input
-	await Promise.all(
-		modelConfigurationIds.map(async (id) => {
-			const { modelConfiguration, modelOptions } = await setupModelInput(id);
-			if (modelConfiguration) allModelConfigurations.value.push(modelConfiguration);
-			if (modelOptions) allModelOptions.value.push(modelOptions);
-		})
-	);
+	const configs = await fetchModelConfigurations(props.node.inputs);
+	if (!configs) return;
+	allModelConfigurations.value = configs.allModelConfigurations;
+	allModelOptions.value = configs.allModelOptions;
 
 	// dataset input
 	if (datasetId.value) {
diff --git a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-node-ciemss.vue b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-node-ciemss.vue
index 571218e601b..0c9ffc7fd41 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-node-ciemss.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ensemble-ciemss/tera-calibrate-ensemble-node-ciemss.vue
@@ -5,6 +5,14 @@
 			:are-embed-actions-visible="false"
 			:visualization-spec="lossChartSpec"
 		/>
+		<div v-if="outputData">
+			<vega-chart
+				v-for="(chart, index) of ensembleVariableCharts"
+				:key="index"
+				:interactive="false"
+				:visualization-spec="chart"
+			/>
+		</div>
 		<tera-progress-spinner
 			v-if="inProgressCalibrationId || inProgressForecastId"
 			:font-size="2"
@@ -22,29 +30,31 @@
 </template>
 
 <script setup lang="ts">
-import { computed, ref, shallowRef, watch, onMounted } from 'vue';
+import { computed, ref, shallowRef, watch, onMounted, toRef } from 'vue';
 import _ from 'lodash';
 import Button from 'primevue/button';
 import TeraOperatorPlaceholder from '@/components/operator/tera-operator-placeholder.vue';
 import TeraProgressSpinner from '@/components/widgets/tera-progress-spinner.vue';
 import {
-	getRunResultCiemss,
 	pollAction,
 	getCalibrateBlobURL,
 	makeEnsembleCiemssSimulation,
-	getSimulation
+	getSimulation,
+	DataArray
 } from '@/services/models/simulation-service';
-import { setupCsvAsset } from '@/services/calibrate-workflow';
+import { parseCsvAsset, setupCsvAsset } from '@/services/calibrate-workflow';
 import { nodeMetadata, nodeOutputLabel } from '@/components/workflow/util';
 import { logger } from '@/utils/logger';
 import { Poller, PollerState } from '@/api/api';
 import type { WorkflowNode } from '@/types/workflow';
 import { WorkflowPortStatus } from '@/types/workflow';
-import type { CsvAsset, EnsembleSimulationCiemssRequest, Dataset, Simulation } from '@/types/Types';
-import type { RunResults } from '@/types/SimulateConfig';
+import type { CsvAsset, EnsembleSimulationCiemssRequest, Dataset, Simulation, ModelConfiguration } from '@/types/Types';
 import { createDatasetFromSimulationResult, getDataset } from '@/services/dataset';
 import VegaChart from '@/components/widgets/VegaChart.vue';
 import { useProjects } from '@/composables/project';
+import { useChartSettings } from '@/composables/useChartSettings';
+import { useCharts } from '@/composables/useCharts';
+import { GroupedDataArray } from '@/services/charts';
 import {
 	CalibrateEnsembleCiemssOperation,
 	CalibrateEnsembleCiemssOperationState
@@ -52,7 +62,11 @@ import {
 import {
 	updateLossChartSpec,
 	getLossValuesFromSimulation,
-	formatCalibrateModelConfigurations
+	formatCalibrateModelConfigurations,
+	getSelectedOutputEnsembleMapping,
+	buildChartData,
+	fetchModelConfigurations,
+	fetchOutputData
 } from './calibrate-ensemble-util';
 
 const props = defineProps<{
@@ -60,9 +74,10 @@ const props = defineProps<{
 }>();
 const emit = defineEmits(['open-drilldown', 'update-state', 'append-output', 'append-input-port']);
 
-const runResults = ref<RunResults>({});
 const csvAsset = shallowRef<CsvAsset | undefined>(undefined);
 
+const allModelConfigurations = ref<ModelConfiguration[]>([]);
+
 const areInputsFilled = computed(
 	() =>
 		props.node.inputs[0].value &&
@@ -72,7 +87,40 @@ const inProgressCalibrationId = computed(() => props.node.state.inProgressCalibr
 const inProgressForecastId = computed(() => props.node.state.inProgressForecastId);
 const lossValues = ref<{ [key: string]: number }[]>([]);
 const lossChartSpec = ref();
-const lossChartSize = { width: 180, height: 120 };
+const chartSize = { width: 180, height: 120 };
+
+// Charts setup
+const outputData = ref<{
+	result: DataArray;
+	resultSummary: DataArray;
+	pyciemssMap: Record<string, string>;
+	resultGroupByTimepoint: GroupedDataArray;
+} | null>(null);
+const groundTruthData = computed<DataArray>(() => parseCsvAsset(csvAsset.value as CsvAsset));
+const selectedOutputMapping = computed(() => getSelectedOutputEnsembleMapping(props.node));
+const { selectedEnsembleVariableSettings } = useChartSettings(props, emit);
+const { useEnsembleVariableCharts } = useCharts(
+	props.node.id,
+	null,
+	allModelConfigurations,
+	computed(() => buildChartData(outputData.value, selectedOutputMapping.value)),
+	toRef(chartSize),
+	null,
+	selectedOutputMapping
+);
+const ensembleVariableCharts = computed(() => {
+	const charts = useEnsembleVariableCharts(selectedEnsembleVariableSettings, groundTruthData);
+	const ensembleCharts = selectedEnsembleVariableSettings.value.map((setting) => {
+		// Grab the first chart only since the rest of the charts are for model configurations charts
+		const spec = charts.value[setting.id][0];
+		// Make sure the chart since width of the chart can be too small if charts were small multiple charts.
+		spec.width = chartSize.width;
+		spec.height = chartSize.height + 100;
+		return spec;
+	});
+	return ensembleCharts;
+});
+// ----------------------------
 
 const poller = new Poller();
 const pollResult = async (runId: string) => {
@@ -86,7 +134,7 @@ const pollResult = async (runId: string) => {
 						iter: i,
 						loss: d.data.loss
 					}));
-				lossChartSpec.value = updateLossChartSpec(lossValues.value, lossChartSize);
+				lossChartSpec.value = updateLossChartSpec(lossValues.value, chartSize);
 			}
 			if (runId === props.node.state.inProgressCalibrationId && data.updates.length > 0) {
 				const checkpoint = _.last(data.updates);
@@ -133,7 +181,14 @@ const pollResult = async (runId: string) => {
 // Init loss chart
 onMounted(async () => {
 	lossValues.value = await getLossValuesFromSimulation(props.node.state.calibrationId);
-	lossChartSpec.value = await updateLossChartSpec(lossValues.value, lossChartSize);
+	lossChartSpec.value = await updateLossChartSpec(lossValues.value, chartSize);
+});
+
+// Fetch model configurations
+onMounted(async () => {
+	const configs = await fetchModelConfigurations(props.node.inputs);
+	if (!configs) return;
+	allModelConfigurations.value = configs.allModelConfigurations;
 });
 
 watch(
@@ -235,12 +290,8 @@ watch(
 		const state = props.node.state;
 		if (!active) return;
 		if (!state.postForecastId) return;
-
-		const forecastId = state.postForecastId;
-
-		// Simulate
-		const result = await getRunResultCiemss(forecastId, 'result.csv');
-		runResults.value = result.runResults;
+		// Fetch output data and prepare chart data
+		outputData.value = await fetchOutputData(state.preForecastId, state.postForecastId);
 
 		// Dataset used to calibrate
 		const datasetId = props.node.inputs[0]?.value?.[0];