diff --git a/packages/client/hmi-client/src/components/models/tera-stratified-model-configuration.vue b/packages/client/hmi-client/src/components/models/tera-stratified-model-configuration.vue
index 7f69f512a4..0f439743c9 100644
--- a/packages/client/hmi-client/src/components/models/tera-stratified-model-configuration.vue
+++ b/packages/client/hmi-client/src/components/models/tera-stratified-model-configuration.vue
@@ -140,6 +140,7 @@ import {
 import { getModelConfigurations } from '@/services/model';
 import TeraStratifiedValueMatrix from '@/components/models/tera-stratified-value-matrix.vue';
 import { NodeType } from '@/model-representation/petrinet/petrinet-renderer';
+import { getBaseAMR } from '@/model-representation/petrinet/petrinet-service';
 import { FeatureConfig } from '@/types/common';
 
 const props = defineProps<{
@@ -162,12 +163,9 @@ const configurations = computed<Model[]>(
 	() => modelConfigurations.value?.map((m) => m.configuration) ?? []
 );
 
-const baseModelStates = computed<any>(() =>
-	props.model?.semantics?.span?.[0].system.model.states.map(({ id }) => id)
-);
-const baseModelTransitions = computed<any>(() =>
-	props.model?.semantics?.span?.[0].system.model.transitions.map(({ id }) => id)
-);
+const baseModel = computed<any>(() => getBaseAMR(props.model));
+const baseModelStates = computed<any>(() => baseModel.value.states.map(({ id }) => id));
+const baseModelTransitions = computed<any>(() => baseModel.value.transitions.map(({ id }) => id));
 
 // Decide if we should display the whole configuration table
 const isConfigurationVisible = computed(
diff --git a/packages/client/hmi-client/src/components/models/tera-stratified-value-matrix.vue b/packages/client/hmi-client/src/components/models/tera-stratified-value-matrix.vue
index 20dc765d03..ac91d1590f 100644
--- a/packages/client/hmi-client/src/components/models/tera-stratified-value-matrix.vue
+++ b/packages/client/hmi-client/src/components/models/tera-stratified-value-matrix.vue
@@ -44,23 +44,23 @@
 							:key="j"
 							@click="
 								valueToEdit = {
-									val: getMatrixValue(cell?.value?.[chosenCol]),
+									val: getMatrixValue(cell?.value?.id),
 									rowIdx: i,
 									colIdx: j
 								}
 							"
 						>
-							<template v-if="cell?.value?.[chosenCol] && cell?.value?.[chosenRow]">
+							<template v-if="cell?.value?.id">
 								<InputText
 									v-if="valueToEdit.rowIdx === i && valueToEdit.colIdx === j"
 									class="cell-input"
 									v-model.lazy="valueToEdit.val"
 									v-focus
 									@focusout="valueToEdit = { val: '', rowIdx: -1, colIdx: -1 }"
-									@keyup.stop.enter="updateModelConfigValue(cell?.value?.[chosenRow])"
+									@keyup.stop.enter="updateModelConfigValue(cell?.value?.id)"
 								/>
 								<span v-else class="editable-cell">
-									{{ getMatrixValue(cell?.value?.[chosenCol]) }}
+									{{ getMatrixValue(cell?.value?.id) }}
 								</span>
 							</template>
 							<span v-else class="not-allowed">N/A</span>
@@ -74,10 +74,10 @@
 
 <script setup lang="ts">
 import { ref, onMounted } from 'vue';
-import { cloneDeep } from 'lodash';
+import { cloneDeep, isEmpty } from 'lodash';
 import {
-	extractStateMatrixData,
-	extractTransitionMatrixData
+	getCatlabStatesMatrixData,
+	getCatlabTransitionsMatrixData
 } from '@/model-representation/petrinet/petrinet-service';
 import { createMatrix1D, createMatrix2D } from '@/utils/pivot';
 import Dropdown from 'primevue/dropdown';
@@ -157,44 +157,31 @@ function updateModelConfigValue(variableName: string) {
 }
 
 function configureMatrix() {
-	// Get stratified ids from the chosen base id
-	const rowAndCol = props.modelConfiguration.configuration?.semantics?.span?.[0]?.map
-		.filter((id: string) => id[1] === props.id)
-		.map((d) => d[0]);
-
-	// Assign dimensions relevant to the ids
-	const rowAndColDimensions = props.modelConfiguration.configuration.semantics?.typing?.map.filter(
-		(id) => rowAndCol.includes(id[0])
-	);
-
-	if (rowAndColDimensions) {
-		// Assuming column id is the first one
-		const colId = rowAndColDimensions[0][0];
-		let rowId: string | null = null;
-
-		for (let i = 0; i < rowAndColDimensions?.length; i++) {
-			// If other id is recognized it is the row id
-			if (!rowId && rowAndColDimensions[i][0] !== colId) {
-				rowId = rowAndColDimensions[i][0];
-			}
-
-			// Push dimension
-			if (colId === rowAndColDimensions[i][0]) {
-				colDimensions.push(rowAndColDimensions[i][1]);
-			} else if (rowId === rowAndColDimensions[i][0]) {
-				rowDimensions.push(rowAndColDimensions[i][1]);
-			}
-		}
-	}
+	console.log(props.modelConfiguration.configuration);
 
 	// Get only the states/transitions that are mapped to the base model
 	const matrixData =
 		props.nodeType === NodeType.State
-			? extractStateMatrixData(props.modelConfiguration.configuration, rowAndCol, [
-					...colDimensions,
-					...rowDimensions
-			  ])
-			: extractTransitionMatrixData(props.modelConfiguration.configuration, rowAndCol);
+			? getCatlabStatesMatrixData(props.modelConfiguration.configuration).filter(
+					(d) => d['@base'] === props.id
+			  )
+			: getCatlabTransitionsMatrixData(props.modelConfiguration.configuration).filter(
+					(d) => d['@base'] === props.id
+			  );
+
+	console.log(matrixData);
+
+	if (isEmpty(matrixData)) return;
+
+	// Grab dimension names from the first matrix row
+	const dimensions = [cloneDeep(matrixData)[0]].map((d) => {
+		delete d.id;
+		delete d['@base'];
+		return Object.keys(d);
+	})[0];
+
+	rowDimensions.push(...dimensions);
+	colDimensions.push(...dimensions);
 
 	const matrixAttributes =
 		props.nodeType === NodeType.State
@@ -204,6 +191,8 @@ function configureMatrix() {
 	matrix.value = matrixAttributes.matrix;
 	chosenCol.value = colDimensions[0];
 	chosenRow.value = rowDimensions[0];
+
+	console.log(matrix.value);
 }
 
 onMounted(() => {
@@ -249,6 +238,7 @@ onMounted(() => {
 
 .cell-input {
 	height: 4rem;
+	min-width: 8rem;
 	width: 100%;
 	padding-left: 12px;
 }
diff --git a/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts b/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts
index 1079c5364f..c34dad840d 100644
--- a/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts
+++ b/packages/client/hmi-client/src/model-representation/petrinet/petrinet-service.ts
@@ -665,6 +665,25 @@ export const extractStateMatrixData = (amr: Model, stateIds: string[], dimension
 
 const MAX_DEPTH = 10;
 
+/**
+ * Extracts the base model used in a stratified model generated by catlab.
+ * Useful if we want to view values of a stratified model from the percpective of the base.
+ */
+export const getBaseAMR = (amr: Model) => {
+	let level = amr?.semantics?.span?.[0].system;
+	let c = 0;
+
+	while (c < MAX_DEPTH) {
+		c++;
+		if (level.semantics.span) {
+			level = level.semantics.span[0].system;
+		} else {
+			return level.model;
+		}
+	}
+	return level.model;
+};
+
 /**
  * Given an identifier string id, recursiviely lookup mapping information
  * to find out where the id originated. At eash iteration record the term