From 37e596f172f83a2b9f941e307c5b9113186fce2e Mon Sep 17 00:00:00 2001
From: Daniel Chang <mwdchang@gmail.com>
Date: Mon, 29 Jul 2024 16:44:22 -0400
Subject: [PATCH] chart tweaks (#4306)

---
 .../tera-calibrate-ciemss-drilldown.vue       |  2 +-
 .../tera-calibrate-node-ciemss.vue            | 13 +++----
 .../tera-optimize-ciemss-drilldown.vue        | 14 +++++---
 .../tera-optimize-ciemss-node.vue             | 20 ++++++++---
 .../simulate-ciemss/tera-simulate-ciemss.vue  |  2 +-
 .../tera-simulate-node-ciemss.vue             | 12 ++++---
 .../service/data/SimulationService.java       | 33 +++++++++--------
 .../notification/NotificationService.java     |  6 ++--
 .../service/data/SimulationServiceTests.java  | 36 +++++++++----------
 9 files changed, 78 insertions(+), 60 deletions(-)

diff --git a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-ciemss-drilldown.vue b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-ciemss-drilldown.vue
index 9608cd277a..61004aac84 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-ciemss-drilldown.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-ciemss-drilldown.vue
@@ -306,7 +306,7 @@ const preparedCharts = computed(() => {
 				groupField: 'sample_id'
 			},
 			{
-				title: `${config.join(',')}`,
+				title: '',
 				width: chartSize.value.width,
 				height: chartSize.value.height,
 				legend: true,
diff --git a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-node-ciemss.vue b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-node-ciemss.vue
index ca6a2ce580..6a503255f3 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-node-ciemss.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/calibrate-ciemss/tera-calibrate-node-ciemss.vue
@@ -41,7 +41,7 @@ import {
 	DataArray
 } from '@/services/models/simulation-service';
 import { getModelConfigurationById, createModelConfiguration } from '@/services/model-configurations';
-import { getModelByModelConfigurationId } from '@/services/model';
+import { getModelByModelConfigurationId, getUnitsFromModelParts } from '@/services/model';
 import { renderLossGraph, setupDatasetInput } from '@/services/calibrate-workflow';
 import { nodeMetadata, nodeOutputLabel } from '@/components/workflow/util';
 import { logger } from '@/utils/logger';
@@ -62,6 +62,8 @@ const emit = defineEmits(['open-drilldown', 'update-state', 'append-output']);
 const modelConfigId = computed<string | undefined>(() => props.node.inputs[0].value?.[0]);
 
 const model = ref<Model | null>(null);
+const modelVarUnits = ref<{ [key: string]: string }>({});
+
 const runResult = ref<DataArray>([]);
 const runResultPre = ref<DataArray>([]);
 const runResultSummary = ref<DataArray>([]);
@@ -142,8 +144,6 @@ const preparedCharts = computed(() => {
 			}
 		});
 
-		const xAxisTitle = model.value?.semantics?.ode.time?.units?.expression ?? 'time';
-
 		return createForecastChart(
 			{
 				dataset: result,
@@ -162,13 +162,13 @@ const preparedCharts = computed(() => {
 				timeField: datasetTimeField as string
 			},
 			{
-				title: `${config.join(',')}`,
+				title: '',
 				width: 180,
 				height: 120,
 				legend: true,
 				translationMap: reverseMap,
-				xAxisTitle,
-				yAxisTitle: `${config.join(',')}`,
+				xAxisTitle: modelVarUnits.value._time || 'Time',
+				yAxisTitle: _.uniq(config.map((v) => modelVarUnits.value[v]).filter((v) => !!v)).join(',') || '',
 				colorscheme: ['#AAB3C6', '#1B8073']
 			}
 		);
@@ -237,6 +237,7 @@ watch(
 
 		const id = input.value[0];
 		model.value = await getModelByModelConfigurationId(id);
+		modelVarUnits.value = getUnitsFromModelParts(model.value as Model);
 	},
 	{ immediate: true }
 );
diff --git a/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-drilldown.vue b/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-drilldown.vue
index ad81d0adc7..060408a6d1 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-drilldown.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-drilldown.vue
@@ -894,7 +894,10 @@ const preparedForecastCharts = computed(() => {
 
 	// intervention chart spec
 	charts.interventionCharts = knobs.value.selectedInterventionVariables.map((variable) => {
-		chartOptions.translationMap = translationMap(variable);
+		const options = _.cloneDeep(chartOptions);
+		options.translationMap = translationMap(variable);
+		options.yAxisTitle = getUnit(variable);
+
 		const forecastChart = createForecastChart(
 			{
 				dataset: result,
@@ -908,7 +911,7 @@ const preparedForecastCharts = computed(() => {
 				timeField: 'timepoint_id'
 			},
 			null,
-			chartOptions
+			options
 		);
 		// add intervention annotations (rules and text)
 		forecastChart.layer.push(...createInterventionChartMarkers(preProcessedInterventionsData.value[variable]));
@@ -917,7 +920,10 @@ const preparedForecastCharts = computed(() => {
 
 	// simulation chart spec
 	charts.simulationCharts = knobs.value.selectedSimulationVariables.map((variable) => {
-		chartOptions.translationMap = translationMap(variable);
+		const options = _.cloneDeep(chartOptions);
+		options.translationMap = translationMap(variable);
+		options.yAxisTitle = getUnit(variable);
+
 		return createForecastChart(
 			{
 				dataset: result,
@@ -931,7 +937,7 @@ const preparedForecastCharts = computed(() => {
 				timeField: 'timepoint_id'
 			},
 			null,
-			chartOptions
+			options
 		);
 	});
 	return charts;
diff --git a/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-node.vue b/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-node.vue
index 2d5fcb6362..b4efc536a4 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-node.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/optimize-ciemss/tera-optimize-ciemss-node.vue
@@ -1,6 +1,6 @@
 <template>
 	<main>
-		<tera-operator-placeholder v-if="!showSpinner" :node="node">
+		<tera-operator-placeholder v-if="!showSpinner && !runResults[node.state.postForecastRunId]" :node="node">
 			<template v-if="!node.inputs[0].value"> Attach a model configuration </template>
 		</tera-operator-placeholder>
 		<template v-if="node.inputs[0].value">
@@ -38,7 +38,9 @@ import { createLLMSummary } from '@/services/summary-service';
 import VegaChart from '@/components/widgets/VegaChart.vue';
 import { createForecastChart } from '@/services/charts';
 import { mergeResults, renameFnGenerator } from '@/components/workflow/ops/calibrate-ciemss/calibrate-utils';
+import { getModelByModelConfigurationId, getUnitsFromModelParts } from '@/services/model';
 import { getModelConfigurationById } from '@/services/model-configurations';
+
 import {
 	OptimizeCiemssOperationState,
 	OptimizeCiemssOperation,
@@ -55,6 +57,8 @@ const runResults = ref<any>({});
 const runResultsSummary = ref<any>({});
 const modelConfigId = computed<string | undefined>(() => props.node.inputs[0]?.value?.[0]);
 
+const modelVarUnits = ref<{ [key: string]: string }>({});
+
 let pyciemssMap: Record<string, string> = {};
 
 const showSpinner = computed<boolean>(
@@ -142,13 +146,16 @@ const preparedCharts = computed(() => {
 				width: 180,
 				height: 120,
 				legend: true,
-				xAxisTitle: 'Time',
-				yAxisTitle: variable,
+				xAxisTitle: modelVarUnits.value._time || 'Time',
+				yAxisTitle:
+					_.uniq(modelVarUnits.value[variable])
+						.filter((v) => !!v)
+						.join(',') || '',
 				translationMap: {
 					[`${pyciemssMap[variable]}_mean:pre`]: `${variable} before optimization`,
 					[`${pyciemssMap[variable]}_mean`]: `${variable} after optimization`
 				},
-				title: variable,
+				title: '',
 				colorscheme: ['#AAB3C6', '#1B8073']
 			}
 		)
@@ -237,6 +244,11 @@ watch(
 		if (!active) return;
 		if (!state.postForecastRunId || !state.preForecastRunId) return;
 
+		const model = await getModelByModelConfigurationId(modelConfigId.value as string);
+		if (model) {
+			modelVarUnits.value = getUnitsFromModelParts(model);
+		}
+
 		const preForecastRunId = state.preForecastRunId;
 		const postForecastRunId = state.postForecastRunId;
 
diff --git a/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-ciemss.vue b/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-ciemss.vue
index de7cb84886..6a7caff5c5 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-ciemss.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-ciemss.vue
@@ -348,7 +348,7 @@ const preparedCharts = computed(() => {
 			null,
 			// options
 			{
-				title: `${config.join(',')}`,
+				title: '',
 				width: chartSize.value.width,
 				height: chartSize.value.height,
 				legend: true,
diff --git a/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-node-ciemss.vue b/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-node-ciemss.vue
index 5fe6f9c5f1..591c7226e5 100644
--- a/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-node-ciemss.vue
+++ b/packages/client/hmi-client/src/components/workflow/ops/simulate-ciemss/tera-simulate-node-ciemss.vue
@@ -27,7 +27,7 @@ import {
 	parsePyCiemssMap,
 	DataArray
 } from '@/services/models/simulation-service';
-import { getModelByModelConfigurationId } from '@/services/model';
+import { getModelByModelConfigurationId, getUnitsFromModelParts } from '@/services/model';
 import { Poller, PollerState } from '@/api/api';
 import { logger } from '@/utils/logger';
 import { chartActionsProxy, nodeOutputLabel } from '@/components/workflow/util';
@@ -45,6 +45,8 @@ const props = defineProps<{
 
 const emit = defineEmits(['open-drilldown', 'update-state', 'append-output']);
 const model = ref<Model | null>(null);
+const modelVarUnits = ref<{ [key: string]: string }>({});
+
 const runResults = ref<{ [runId: string]: DataArray }>({});
 const runResultsSummary = ref<{ [runId: string]: DataArray }>({});
 
@@ -147,7 +149,6 @@ const preparedCharts = computed(() => {
 	Object.keys(pyciemssMap).forEach((key) => {
 		reverseMap[`${pyciemssMap[key]}_mean`] = key;
 	});
-	const xAxisTitle = model.value?.semantics?.ode.time?.units?.expression ?? 'time';
 
 	return props.node.state.chartConfigs.map((config) =>
 		createForecastChart(
@@ -165,13 +166,13 @@ const preparedCharts = computed(() => {
 			null,
 			// options
 			{
-				title: `${config.join(',')}`,
+				title: '',
 				width: 180,
 				height: 120,
 				legend: true,
 				translationMap: reverseMap,
-				xAxisTitle,
-				yAxisTitle: `${config.join(',')}`
+				xAxisTitle: modelVarUnits.value._time || 'Time',
+				yAxisTitle: _.uniq(config.map((v) => modelVarUnits.value[v]).filter((v) => !!v)).join(',') || ''
 			}
 		)
 	);
@@ -185,6 +186,7 @@ watch(
 
 		const id = input.value[0];
 		model.value = await getModelByModelConfigurationId(id);
+		modelVarUnits.value = getUnitsFromModelParts(model.value as Model);
 	},
 	{ immediate: true }
 );
diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/SimulationService.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/SimulationService.java
index f4ffdde352..441e503c67 100644
--- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/SimulationService.java
+++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/data/SimulationService.java
@@ -1,12 +1,9 @@
 package software.uncharted.terarium.hmiserver.service.data;
 
-import java.util.UUID;
-
-import org.springframework.stereotype.Service;
-
 import com.fasterxml.jackson.databind.ObjectMapper;
-
 import io.micrometer.observation.annotation.Observed;
+import java.util.UUID;
+import org.springframework.stereotype.Service;
 import software.uncharted.terarium.hmiserver.configuration.Config;
 import software.uncharted.terarium.hmiserver.models.dataservice.dataset.Dataset;
 import software.uncharted.terarium.hmiserver.models.dataservice.simulation.Simulation;
@@ -31,13 +28,14 @@ public class SimulationService extends TerariumAssetServiceWithoutSearch<Simulat
 	 * @param s3ClientService     S3 client service
 	 */
 	public SimulationService(
-			final ObjectMapper objectMapper,
-			final Config config,
-			final ProjectService projectService,
-			final ProjectAssetService projectAssetService,
-			final SimulationRepository repository,
-			final S3ClientService s3ClientService,
-			final SimulationUpdateRepository simulationUpdateRepository) {
+		final ObjectMapper objectMapper,
+		final Config config,
+		final ProjectService projectService,
+		final ProjectAssetService projectAssetService,
+		final SimulationRepository repository,
+		final S3ClientService s3ClientService,
+		final SimulationUpdateRepository simulationUpdateRepository
+	) {
 		super(objectMapper, config, projectService, projectAssetService, repository, s3ClientService, Simulation.class);
 		this.simulationUpdateRepository = simulationUpdateRepository;
 	}
@@ -56,9 +54,10 @@ private String getDatasetPath(final UUID datasetId, final String filename) {
 	}
 
 	public SimulationUpdate appendUpdateToSimulation(
-			final UUID simulationId,
-			final SimulationUpdate update,
-			final Schema.Permission hasReadPermission) {
+		final UUID simulationId,
+		final SimulationUpdate update,
+		final Schema.Permission hasReadPermission
+	) {
 		final Simulation simulation = getAsset(simulationId, hasReadPermission).orElseThrow();
 
 		final SimulationUpdate updated = simulationUpdateRepository.save(update);
@@ -80,8 +79,8 @@ public void copySimulationResultToDataset(final Simulation simulation, final Dat
 				final String srcPath = getResultsPath(simId, filename);
 				final String destPath = getDatasetPath(dataset.getId(), filename);
 				s3ClientService
-						.getS3Service()
-						.copyObject(config.getFileStorageS3BucketName(), srcPath, config.getFileStorageS3BucketName(), destPath);
+					.getS3Service()
+					.copyObject(config.getFileStorageS3BucketName(), srcPath, config.getFileStorageS3BucketName(), destPath);
 			}
 		}
 	}
diff --git a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/notification/NotificationService.java b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/notification/NotificationService.java
index 14ef9a0528..396a92d030 100644
--- a/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/notification/NotificationService.java
+++ b/packages/server/src/main/java/software/uncharted/terarium/hmiserver/service/notification/NotificationService.java
@@ -1,16 +1,14 @@
 package software.uncharted.terarium.hmiserver.service.notification;
 
+import io.micrometer.observation.annotation.Observed;
 import java.sql.Timestamp;
 import java.time.Instant;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
-
-import org.springframework.stereotype.Service;
-
-import io.micrometer.observation.annotation.Observed;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
+import org.springframework.stereotype.Service;
 import software.uncharted.terarium.hmiserver.models.User;
 import software.uncharted.terarium.hmiserver.models.notification.NotificationEvent;
 import software.uncharted.terarium.hmiserver.models.notification.NotificationGroup;
diff --git a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/SimulationServiceTests.java b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/SimulationServiceTests.java
index 9ca97f56f4..362cadbe79 100644
--- a/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/SimulationServiceTests.java
+++ b/packages/server/src/test/java/software/uncharted/terarium/hmiserver/service/data/SimulationServiceTests.java
@@ -1,20 +1,17 @@
 package software.uncharted.terarium.hmiserver.service.data;
 
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
 import java.util.UUID;
-
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.security.test.context.support.WithUserDetails;
-
-import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.ObjectMapper;
-
 import software.uncharted.terarium.hmiserver.TerariumApplicationTests;
 import software.uncharted.terarium.hmiserver.configuration.MockUser;
 import software.uncharted.terarium.hmiserver.models.dataservice.project.Project;
@@ -48,7 +45,8 @@ static SimulationRequest createSimRequest() {
 	@BeforeEach
 	public void setup() throws IOException {
 		project = projectService.createProject(
-				(Project) new Project().setPublicAsset(true).setName("test-project-name").setDescription("my description"));
+			(Project) new Project().setPublicAsset(true).setName("test-project-name").setDescription("my description")
+		);
 	}
 
 	static Simulation createSimulation(final String key) {
@@ -115,9 +113,10 @@ public void testItCanGetSimulations() throws IOException {
 	@WithUserDetails(MockUser.URSULA)
 	public void testItCanGetSimulationById() throws IOException {
 		final Simulation simulation = simulationService.createAsset(
-				createSimulation("0"),
-				project.getId(),
-				ASSUME_WRITE_PERMISSION);
+			createSimulation("0"),
+			project.getId(),
+			ASSUME_WRITE_PERMISSION
+		);
 		final Simulation fetchedSimulation = simulationService.getAsset(simulation.getId(), ASSUME_WRITE_PERMISSION).get();
 
 		Assertions.assertEquals(simulation, fetchedSimulation);
@@ -133,14 +132,15 @@ public void testItCanGetSimulationById() throws IOException {
 	@WithUserDetails(MockUser.URSULA)
 	public void testItCanUpdateSimulation() throws Exception {
 		final Simulation simulation = simulationService.createAsset(
-				createSimulation("A"),
-				project.getId(),
-				ASSUME_WRITE_PERMISSION);
+			createSimulation("A"),
+			project.getId(),
+			ASSUME_WRITE_PERMISSION
+		);
 		simulation.setName("new name");
 
 		final Simulation updatedSimulation = simulationService
-				.updateAsset(simulation, project.getId(), ASSUME_WRITE_PERMISSION)
-				.orElseThrow();
+			.updateAsset(simulation, project.getId(), ASSUME_WRITE_PERMISSION)
+			.orElseThrow();
 
 		Assertions.assertEquals(simulation, updatedSimulation);
 		Assertions.assertNotNull(updatedSimulation.getUpdatedOn());
@@ -150,9 +150,10 @@ public void testItCanUpdateSimulation() throws Exception {
 	@WithUserDetails(MockUser.URSULA)
 	public void testItCanDeleteSimulation() throws Exception {
 		final Simulation simulation = simulationService.createAsset(
-				createSimulation("B"),
-				project.getId(),
-				ASSUME_WRITE_PERMISSION);
+			createSimulation("B"),
+			project.getId(),
+			ASSUME_WRITE_PERMISSION
+		);
 
 		simulationService.deleteAsset(simulation.getId(), project.getId(), ASSUME_WRITE_PERMISSION);
 
@@ -204,7 +205,6 @@ public void testItCanCreateSimulationUpdates() {
 
 			final Simulation again = simulationService.getAsset(after.getId(), ASSUME_WRITE_PERMISSION).get();
 			Assertions.assertEquals(2, again.getUpdates().size());
-
 		} catch (final Exception e) {
 			Assertions.fail(e);
 		}