Add from_dataframe method to Measurement to create multidim measurement from a DataFrame.

This also adds some symmetry with the existing to_dataframe method.

PiperOrigin-RevId: 721549130

Author: glados-verma

openhtf/core/measurements.py

@@ -480,6 +480,34 @@ def to_dataframe(self, columns: Any = None) -> Any:
 
     return dataframe
 
+  def from_dataframe(self, dataframe: Any, metric_column: str):
+    """Convert a pandas DataFrame to a multi-dim measurement.
+
+    Args:
+      dataframe: A pandas DataFrame. Dimensions for this multi-dim measurement
+        need to match columns in the DataFrame (can be multi-index).
+      metric_column: The column name of the metric to be measured.
+
+    Raises:
+      ValueError: If dataframe is missing dimensions.
+    """
+    if not isinstance(self._measured_value, DimensionedMeasuredValue):
+      raise TypeError(
+          'Only a dimensioned measurement can be set from a DataFrame'
+      )
+    dimension_labels = [d.name for d in self.dimensions]
+    dimensioned_df = dataframe.reset_index()
+    try:
+      dimensioned_df.set_index(dimension_labels, inplace=True)
+    except KeyError as e:
+      raise ValueError('DataFrame is missing dimensions') from e
+    if metric_column not in dimensioned_df.columns:
+      raise ValueError(
+          f'DataFrame does not have a column named {metric_column}'
+      )
+    for row_dimensions, row_metrics in dimensioned_df.iterrows():
+      self.measured_value[row_dimensions] = row_metrics[metric_column]
+
 
 @attr.s(slots=True)
 class MeasuredValue(object):

test/core/measurements_test.py

@@ -25,6 +25,7 @@
 from openhtf.core import measurements
 from examples import all_the_things
 from openhtf.util import test as htf_test
+import pandas
 
 # Fields that are considered 'volatile' for record comparison.
 _VOLATILE_FIELDS = {
@@ -105,18 +106,24 @@ def test_precision(self):
     m.measured_value.set(1.2346)
     self.assertAlmostEqual(m.measured_value.value, 1.235)
 
-    m = htf.Measurement('meas_with_precision_and_dims').with_precision(
-        3).with_dimensions('x')
+    m = (
+        htf.Measurement('meas_with_precision_and_dims')
+        .with_precision(3)
+        .with_dimensions('x')
+    )
     m.measured_value[42] = 1.2346
     self.assertAlmostEqual(m.measured_value[42], 1.235)
 
   def test_cache_same_object(self):
     m = htf.Measurement('measurement')
     basetypes0 = m.as_base_types()
-    self.assertEqual({
-        'name': 'measurement',
-        'outcome': 'UNSET',
-    }, basetypes0)
+    self.assertEqual(
+        {
+            'name': 'measurement',
+            'outcome': 'UNSET',
+        },
+        basetypes0,
+    )
     basetypes1 = m.as_base_types()
     self.assertIs(basetypes0, basetypes1)
     m.measured_value.set(1)
@@ -127,7 +134,9 @@ def test_cache_same_object(self):
             'name': 'measurement',
             'outcome': 'PASS',
             'measured_value': 1,
-        }, basetypes2)
+        },
+        basetypes2,
+    )
     self.assertIs(basetypes0, basetypes2)
 
   @htf_test.patch_plugs(user_mock='openhtf.plugs.user_input.UserInput')
@@ -142,52 +151,66 @@ def test_chaining_in_measurement_declarations(self, user_mock):
   @htf_test.yields_phases
   def test_measurements_with_dimensions(self):
     record = yield all_the_things.dimensions
-    self.assertMeasured(record, 'dimensions', [
-        (0, 1),
-        (1, 2),
-        (2, 4),
-        (3, 8),
-        (4, 16),
-    ])
-    self.assertMeasured(record, 'lots_of_dims', [
-        (1, 21, 101, 123),
-        (2, 22, 102, 126),
-        (3, 23, 103, 129),
-        (4, 24, 104, 132),
-    ])
+    self.assertMeasured(
+        record,
+        'dimensions',
+        [
+            (0, 1),
+            (1, 2),
+            (2, 4),
+            (3, 8),
+            (4, 16),
+        ],
+    )
+    self.assertMeasured(
+        record,
+        'lots_of_dims',
+        [
+            (1, 21, 101, 123),
+            (2, 22, 102, 126),
+            (3, 23, 103, 129),
+            (4, 24, 104, 132),
+        ],
+    )
 
   @htf_test.yields_phases
   def test_validator_replacement(self):
     record = yield all_the_things.measures_with_args.with_args(
-        minimum=2, maximum=4)
+        minimum=2, maximum=4
+    )
     self.assertMeasurementFail(record, 'replaced_min_only')
     self.assertMeasurementPass(record, 'replaced_max_only')
     self.assertMeasurementFail(record, 'replaced_min_max')
     record = yield all_the_things.measures_with_args.with_args(
-        minimum=0, maximum=5)
+        minimum=0, maximum=5
+    )
     self.assertMeasurementPass(record, 'replaced_min_only')
     self.assertMeasurementPass(record, 'replaced_max_only')
     self.assertMeasurementPass(record, 'replaced_min_max')
     record = yield all_the_things.measures_with_args.with_args(
-        minimum=-1, maximum=0)
+        minimum=-1, maximum=0
+    )
     self.assertMeasurementPass(record, 'replaced_min_only')
     self.assertMeasurementFail(record, 'replaced_max_only')
     self.assertMeasurementFail(record, 'replaced_min_max')
 
   @htf_test.yields_phases
   def test_validator_replacement_marginal(self):
     record = yield all_the_things.measures_with_marginal_args.with_args(
-        marginal_minimum=4, marginal_maximum=6)
+        marginal_minimum=4, marginal_maximum=6
+    )
     self.assertMeasurementMarginal(record, 'replaced_marginal_min_only')
     self.assertMeasurementNotMarginal(record, 'replaced_marginal_max_only')
     self.assertMeasurementMarginal(record, 'replaced_marginal_min_max')
     record = yield all_the_things.measures_with_marginal_args.with_args(
-        marginal_minimum=1, marginal_maximum=2)
+        marginal_minimum=1, marginal_maximum=2
+    )
     self.assertMeasurementNotMarginal(record, 'replaced_marginal_min_only')
     self.assertMeasurementMarginal(record, 'replaced_marginal_max_only')
     self.assertMeasurementMarginal(record, 'replaced_marginal_min_max')
     record = yield all_the_things.measures_with_marginal_args.with_args(
-        marginal_minimum=2, marginal_maximum=4)
+        marginal_minimum=2, marginal_maximum=4
+    )
     self.assertMeasurementNotMarginal(record, 'replaced_marginal_min_only')
     self.assertMeasurementNotMarginal(record, 'replaced_marginal_max_only')
     self.assertMeasurementNotMarginal(record, 'replaced_marginal_min_max')
@@ -196,12 +219,15 @@ def test_validator_replacement_marginal(self):
   def test_measurement_order(self):
     record = yield all_the_things.dimensions
     self.assertEqual(
-        list(record.measurements.keys()), ['dimensions', 'lots_of_dims'])
+        list(record.measurements.keys()), ['dimensions', 'lots_of_dims']
+    )
     record = yield all_the_things.measures_with_args.with_args(
-        minimum=2, maximum=4)
+        minimum=2, maximum=4
+    )
     self.assertEqual(
         list(record.measurements.keys()),
-        ['replaced_min_only', 'replaced_max_only', 'replaced_min_max'])
+        ['replaced_min_only', 'replaced_max_only', 'replaced_min_max'],
+    )
 
   @htf_test.yields_phases
   def test_bad_validation(self):
@@ -231,14 +257,19 @@ def test_to_dataframe__no_pandas(self):
     with self.assertRaises(RuntimeError):
       self.test_to_dataframe(units=True)
 
-  def test_to_dataframe(self, units=True):
+  def _make_multidim_measurement(self, units=''):
     measurement = htf.Measurement('test_multidim')
     measurement.with_dimensions('ms', 'assembly', htf.Dimension('my_zone'))
+    if units:
+      measurement.with_units(units)
+    return measurement
 
+  def test_to_dataframe(self, units=True):
     if units:
-      measurement.with_units('°C')
+      measurement = self._make_multidim_measurement('°C')
       measure_column_name = 'degree Celsius'
     else:
+      measurement = self._make_multidim_measurement()
       measure_column_name = 'value'
 
     for t in range(5):
@@ -260,6 +291,104 @@ def test_to_dataframe(self, units=True):
   def test_to_dataframe__no_units(self):
     self.test_to_dataframe(units=False)
 
+  def test_from_dataframe_raises_if_dimensions_missing_in_dataframe(self):
+    measurement = self._make_multidim_measurement('°C')
+    with self.assertRaisesRegex(ValueError, 'DataFrame is missing dimensions'):
+      measurement.from_dataframe(
+          pandas.DataFrame({
+              'ms': [1, 2, 3],
+              'my_zone': ['X', 'Y', 'Z'],
+              'degree_celsius': [10, 20, 30],
+          }),
+          metric_column='degree_celsius',
+      )
+
+  def test_from_dataframe_raises_if_metric_missing_in_dataframe(self):
+    measurement = self._make_multidim_measurement('°C')
+    with self.assertRaisesRegex(
+        ValueError, 'DataFrame does not have a column named degree_celsius'
+    ):
+      measurement.from_dataframe(
+          pandas.DataFrame({
+              'ms': [1, 2, 3],
+              'assembly': ['A', 'B', 'C'],
+              'my_zone': ['X', 'Y', 'Z'],
+              'degrees_fahrenheit': [10, 20, 30],
+          }),
+          metric_column='degree_celsius',
+      )
+
+  def test_from_flat_dataframe(self):
+    measurement = self._make_multidim_measurement('°C')
+    source_dataframe = pandas.DataFrame({
+        'ms': [1, 2, 3],
+        'assembly': ['A', 'B', 'C'],
+        'my_zone': ['X', 'Y', 'Z'],
+        'degree_celsius': [10, 20, 30],
+    })
+    measurement.from_dataframe(source_dataframe, metric_column='degree_celsius')
+    measurement.outcome = measurements.Outcome.PASS
+    self.assertEqual(measurement.measured_value[(1, 'A', 'X')], 10)
+    self.assertEqual(measurement.measured_value[(2, 'B', 'Y')], 20)
+    self.assertEqual(measurement.measured_value[(3, 'C', 'Z')], 30)
+    pandas.testing.assert_frame_equal(
+        measurement.to_dataframe().rename(
+            columns={
+                'ms': 'ms',
+                'assembly': 'assembly',
+                'my_zone': 'my_zone',
+                # The metric column name comes from the unit.
+                'degree Celsius': 'degree_celsius',
+            }
+        ),
+        source_dataframe,
+    )
+
+  def test_from_multiindex_dataframe(self):
+    measurement = self._make_multidim_measurement('°C')
+    source_dataframe = pandas.DataFrame({
+        'ms': [1, 2, 3],
+        'assembly': ['A', 'B', 'C'],
+        'my_zone': ['X', 'Y', 'Z'],
+        'degree_celsius': [10, 20, 30],
+    })
+    source_dataframe.set_index(['ms', 'assembly', 'my_zone'], inplace=True)
+    measurement.from_dataframe(source_dataframe, metric_column='degree_celsius')
+    measurement.outcome = measurements.Outcome.PASS
+    self.assertEqual(measurement.measured_value[(1, 'A', 'X')], 10)
+    self.assertEqual(measurement.measured_value[(2, 'B', 'Y')], 20)
+    self.assertEqual(measurement.measured_value[(3, 'C', 'Z')], 30)
+
+  def test_from_dataframe_with_extra_columns(self):
+    measurement = self._make_multidim_measurement('°C')
+    source_dataframe = pandas.DataFrame({
+        'ms': [1, 2, 3],
+        'assembly': ['A', 'B', 'C'],
+        'my_zone': ['X', 'Y', 'Z'],
+        'degree_celsius': [10, 20, 30],
+        'degrees_fahrenheit': [11, 21, 31],
+    })
+    measurement.from_dataframe(source_dataframe, metric_column='degree_celsius')
+    measurement.outcome = measurements.Outcome.PASS
+    self.assertEqual(measurement.measured_value[(1, 'A', 'X')], 10)
+    self.assertEqual(measurement.measured_value[(2, 'B', 'Y')], 20)
+    self.assertEqual(measurement.measured_value[(3, 'C', 'Z')], 30)
+
+  def test_from_dataframe_with_duplicate_dimensions_overwrites(self):
+    measurement = self._make_multidim_measurement('°C')
+    source_dataframe = pandas.DataFrame({
+        'ms': [1, 2, 3, 1],
+        'assembly': ['A', 'B', 'C', 'A'],
+        'my_zone': ['X', 'Y', 'Z', 'X'],
+        'degree_celsius': [10, 20, 30, 11],
+    })
+    measurement.from_dataframe(source_dataframe, metric_column='degree_celsius')
+    measurement.outcome = measurements.Outcome.PASS
+    # Overwritten value.
+    self.assertEqual(measurement.measured_value[(1, 'A', 'X')], 11)
+    self.assertEqual(measurement.measured_value[(2, 'B', 'Y')], 20)
+    self.assertEqual(measurement.measured_value[(3, 'C', 'Z')], 30)
+
   def test_bad_validator(self):
     measurement = htf.Measurement('bad_measure')
     measurement.with_dimensions('a')
@@ -346,8 +475,10 @@ def test_multi_dimension_correct(self):
     try:
       measurement.measured_value[dimension_vals] = 42
     except measurements.InvalidDimensionsError:
-      self.fail('measurement.DimensionedMeasuredValue.__setitem__ '
-                'raised error unexpectedly.')
+      self.fail(
+          'measurement.DimensionedMeasuredValue.__setitem__ '
+          'raised error unexpectedly.'
+      )
 
   def test_multi_dimension_not_enough_error(self):
     measurement = htf.Measurement('measure')