AnalysisCard refactor (#2589)

Summary:
Pull Request resolved: #2589

New structure for Analysis to make the Ax classes easier to reason about in light of storage considerations. The central change is to break up the existing Analysis class into **Analysis** (code for "how" a plot is to be generated) and **AnalysisCard** (the generated plot and its raw data).

See N5607609 for demo.

## Analysis
* Init method to take in "settings" (like which parameter or metrics to operate on, or whether to use observed or modeled effects) and store on class
* compute method takes in Experiment or GenerationStrategy (or both) and outputs an AnalysisCard
* Analysis is a Protocol to allow users to easily implement their own analyses structurally

## AnalysisCard
* Contains the raw data computed by the analysis and a "blob" which holds data processed and ready for end-user consumption (ex. a plot)
* Contains other miscellaneous metadata that can be useful for rendering the card or a collection of cards

See the following diff for sample usage.

Reviewed By: Cesar-Cardoso

Differential Revision: D59926999

fbshipit-source-id: 54c06173753b8d6ba08fe93a9c91a68c48869312

ax/analysis/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from ax.analysis.analysis import Analysis, AnalysisCard, AnalysisCardLevel
+from ax.analysis.markdown import *  # noqa
+from ax.analysis.plotly import *  # noqa
+
+__all__ = ["Analysis", "AnalysisCard", "AnalysisCardLevel"]

ax/analysis/analysis.py

@@ -0,0 +1,87 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from enum import Enum
+from typing import Any, Optional, Protocol
+
+import pandas as pd
+from ax.core.experiment import Experiment
+from ax.modelbridge.generation_strategy import GenerationStrategy
+
+
+class AnalysisCardLevel(Enum):
+    DEBUG = 0
+    LOW = 1
+    MID = 2
+    HIGH = 3
+    CRITICAL = 4
+
+
+class AnalysisCard:
+    # Name of the analysis computed, usually the class name of the Analysis which
+    # produced the card. Useful for grouping by when querying a large collection of
+    # cards.
+    name: str
+
+    title: str
+    subtitle: str
+    level: AnalysisCardLevel
+
+    df: pd.DataFrame  # Raw data produced by the Analysis
+
+    # pyre-ignore[4] We explicitly want to allow any type here, blob is narrowed in
+    # AnalysisCard's subclasses
+    blob: Any  # Data processed and ready for end-user consumption
+
+    # How to interpret the blob (ex. "dataframe", "plotly", "markdown")
+    blob_annotation = "dataframe"
+
+    def __init__(
+        self,
+        name: str,
+        title: str,
+        subtitle: str,
+        level: AnalysisCardLevel,
+        df: pd.DataFrame,
+        # pyre-ignore[2] We explicitly want to allow any type here, blob is narrowed in
+        # AnalysisCard's subclasses
+        blob: Any,
+    ) -> None:
+        self.name = name
+        self.title = title
+        self.subtitle = subtitle
+        self.level = level
+        self.df = df
+        self.blob = blob
+
+
+class Analysis(Protocol):
+    """
+    An Analysis is a class that given either and Experiment, a GenerationStrategy, or
+    both can compute some data intended for end-user consumption. The data is returned
+    to the user in the form of an AnalysisCard which contains the raw data, a blob (the
+    data processed for end-user consumption), and miscellaneous metadata that can be
+    useful for rendering the card or a collection of cards.
+
+    The AnalysisCard is a thin wrapper around the raw data and the processed blob;
+    Analyses impose structure on their blob should subclass Analysis. See
+    PlotlyAnalysis for an example which produces cards where the blob is always a
+    Plotly Figure object.
+
+    A good pattern to follow when implementing your own Analyses is to configure
+    "settings" (like which parameter or metrics to operate on, or whether to use
+    observed or modeled effects) in your Analyses' __init__ methods, then to consume
+    these settings in the compute method.
+    """
+
+    def compute(
+        self,
+        experiment: Optional[Experiment] = None,
+        generation_strategy: Optional[GenerationStrategy] = None,
+    ) -> AnalysisCard:
+        # Note: when implementing compute always prefer experiment.lookup_data() to
+        # experiment.fetch_data() to avoid unintential data fetching within the report
+        # generation.
+        ...

ax/analysis/markdown/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from ax.analysis.markdown.markdown_analysis import (
+    MarkdownAnalysis,
+    MarkdownAnalysisCard,
+)
+
+__all__ = ["MarkdownAnalysis", "MarkdownAnalysisCard"]

ax/analysis/markdown/markdown_analysis.py

@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import pandas as pd
+from ax.analysis.analysis import Analysis, AnalysisCard, AnalysisCardLevel
+from ax.core.experiment import Experiment
+from ax.modelbridge.generation_strategy import GenerationStrategy
+
+
+class MarkdownAnalysisCard(AnalysisCard):
+    name: str
+
+    title: str
+    subtitle: str
+    level: AnalysisCardLevel
+
+    df: pd.DataFrame
+    blob: str
+    blob_annotation = "markdown"
+
+    def get_markdown(self) -> str:
+        return self.blob
+
+
+class MarkdownAnalysis(Analysis):
+    """
+    An Analysis that computes a paragraph of Markdown formatted text.
+    """
+
+    def compute(
+        self,
+        experiment: Optional[Experiment] = None,
+        generation_strategy: Optional[GenerationStrategy] = None,
+    ) -> MarkdownAnalysisCard: ...

ax/analysis/plotly/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from ax.analysis.plotly.plotly_analysis import PlotlyAnalysis, PlotlyAnalysisCard
+
+__all__ = ["PlotlyAnalysis", "PlotlyAnalysisCard"]

ax/analysis/plotly/plotly_analysis.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import pandas as pd
+from ax.analysis.analysis import Analysis, AnalysisCard, AnalysisCardLevel
+from ax.core.experiment import Experiment
+from ax.modelbridge.generation_strategy import GenerationStrategy
+from plotly import graph_objects as go
+
+
+class PlotlyAnalysisCard(AnalysisCard):
+    name: str
+
+    title: str
+    subtitle: str
+    level: AnalysisCardLevel
+
+    df: pd.DataFrame
+    blob: go.Figure
+    blob_annotation = "plotly"
+
+    def get_figure(self) -> go.Figure:
+        return self.blob
+
+
+class PlotlyAnalysis(Analysis):
+    """
+    An Analysis that computes a Plotly figure.
+    """
+
+    def compute(
+        self,
+        experiment: Optional[Experiment] = None,
+        generation_strategy: Optional[GenerationStrategy] = None,
+    ) -> PlotlyAnalysisCard: ...

sphinx/source/analysis.rst

@@ -0,0 +1,32 @@
+.. role:: hidden
+    :class: hidden-section
+
+ax.analysis
+===========
+
+.. automodule:: ax.analysis
+.. currentmodule:: ax.analysis
+
+Analysis
+~~~~~~~~
+
+.. automodule:: ax.analysis.analysis
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+Markdown Analysis
+~~~~~~~~~~~~~~~~~
+
+.. automodule:: ax.analysis.markdown.markdown_analysis
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+Plotly Analysis
+~~~~~~~~~~~~~~~
+
+.. automodule:: ax.analysis.plotly.plotly_analysis
+    :members:
+    :undoc-members:
+    :show-inheritance: