Merge MapData into Data (facebook#4715)

Summary:

This diff merges MapData into Data by giving Data an attribute `has_step_column`. MapData becomes an empty subclass which will be removed in a subsequent PR (D89814417).

Note on how this diff is split up: In this diff, functions that previously required MapData now can consume Data, and type checks stop referencing MapData. However, functions can still return MapData; all references are removed in the next diff.

Changes:
* Functionality from map_data.py is moved into data.py, and functionality from MapData moves into Data; MapData is an empty subclass.
* Data (and MapData) get an attribute `has_step_column`; the important distinction becomes `has_step_column`, not type. * It is now possible for both `Data` and `MapData` to either have a step column or not. Having both `Data` and `MapData` like this is an unpleasant intermediate state that we should move off of immediately (landing this and D89814417 together)
* Many `isinstance` checks become `has_step_coumn` checks
* Data's new methods `subsample` and `latest` get special cases for when there is no "step" column
* Make Data's required columns always be the same and not contain "step" (since it isn't really required). Remove `required_columns` method.

Reviewed By: lena-kashtelyan

Differential Revision: D89820078

Author: esantorella

ax/adapter/data_utils.py

@@ -22,9 +22,8 @@
 from typing import Any
 
 import numpy as np
-from ax.core.data import Data
+from ax.core.data import Data, MAP_KEY
 from ax.core.experiment import Experiment
-from ax.core.map_data import MAP_KEY, MapData
 from ax.core.map_metric import MapMetric
 from ax.core.observation import Observation, ObservationData, ObservationFeatures
 from ax.core.trial_status import STATUSES_EXPECTING_DATA, TrialStatus
@@ -403,7 +402,7 @@ def _extract_observation_data(
         retrived with ``("metadata", "start_time")``.
     """
     data = data if data is not None else experiment.lookup_data()
-    if isinstance(data, MapData):
+    if data.has_step_column:
         if data_loader_config.latest_rows_per_group is not None:
             data = data.latest(
                 rows_per_group=data_loader_config.latest_rows_per_group,
@@ -455,7 +454,7 @@ def _extract_observation_data(
 
     # Identify potential metadata columns.
     index_cols = ["trial_index", "arm_name"]
-    if isinstance(data, MapData):
+    if data.has_step_column:
         index_cols.append(MAP_KEY)
 
     standard_columns = set(index_cols).union(

ax/analysis/overview.py

@@ -31,7 +31,6 @@
 from ax.core.analysis_card import AnalysisCardGroup
 from ax.core.batch_trial import BatchTrial
 from ax.core.experiment import Experiment
-from ax.core.map_data import MapData
 from ax.core.map_metric import MapMetric
 from ax.core.trial_status import TrialStatus
 from ax.exceptions.core import UserInputError
@@ -179,7 +178,7 @@ def compute(
 
         # Check if the experiment has MapData and MapMetrics (required for
         # early stopping)
-        has_map_data = isinstance(experiment.lookup_data(), MapData)
+        has_map_data = experiment.lookup_data().has_step_column
         has_map_metrics = any(
             isinstance(m, MapMetric) for m in experiment.metrics.values()
         )

ax/analysis/plotly/progression.py

@@ -19,7 +19,7 @@
 from ax.analysis.plotly.utils import select_metric
 from ax.analysis.utils import validate_experiment
 from ax.core.experiment import Experiment
-from ax.core.map_data import MAP_KEY, MapData
+from ax.core.map_data import MAP_KEY
 from ax.core.trial_status import TrialStatus
 from ax.generation_strategy.generation_strategy import GenerationStrategy
 
@@ -87,8 +87,8 @@ def validate_applicable_state(
             return experiment_invalid_reason
 
         data = none_throws(experiment).lookup_data()
-        if not isinstance(data, MapData):
-            return "Requires MapData."
+        if not data.has_step_column:
+            return "Requires data to have a column 'step.'"
 
     @override
     def compute(

ax/analysis/results.py

@@ -250,10 +250,9 @@ def compute(
         # Compute progression plots for MapMetrics (learning curves)
         progression_group = None
         data = experiment.lookup_data()
-        has_map_data = isinstance(data, MapData)
         metrics = experiment.metrics.values()
         map_metrics = [m for m in metrics if isinstance(m, MapMetric)]
-        if has_map_data and len(map_metrics) > 0:
+        if data.has_step_column and len(map_metrics) > 0:
             progression_cards = [
                 ProgressionPlot(
                     metric_name=m.name, by_wallclock_time=by_wallclock_time

ax/analysis/summary.py

@@ -14,7 +14,6 @@
 from ax.analysis.utils import validate_experiment
 from ax.core.analysis_card import AnalysisCard
 from ax.core.experiment import Experiment
-from ax.core.map_data import MapData
 from ax.core.trial_status import NON_STALE_STATUSES, TrialStatus
 from ax.exceptions.core import UserInputError
 from ax.generation_strategy.generation_strategy import GenerationStrategy
@@ -87,7 +86,7 @@ def compute(
         should_relativize = (
             len(experiment.metrics) > 0
             and experiment.status_quo is not None
-            and not isinstance(data, MapData)
+            and not data.has_step_column
         )
 
         return self._create_analysis_card(

ax/core/base_trial.py

@@ -20,7 +20,6 @@
 from ax.core.data import Data
 from ax.core.evaluations_to_data import raw_evaluations_to_data
 from ax.core.generator_run import GeneratorRun, GeneratorRunType
-from ax.core.map_data import MapData
 from ax.core.metric import Metric, MetricFetchResult
 from ax.core.runner import Runner
 from ax.core.trial_status import TrialStatus
@@ -449,7 +448,7 @@ def fetch_data(self, metrics: list[Metric] | None = None, **kwargs: Any) -> Data
         data = Metric._unwrap_trial_data_multi(
             results=self.fetch_data_results(metrics=metrics, **kwargs)
         )
-        if not isinstance(data, MapData):
+        if not data.has_step_column:
             data.full_df = sort_by_trial_index_and_arm_name(data.full_df)
 
         return data