Repository: OpenSourceEconomics/respy

Issue:
Add comparison_plot_data to msm().
* respy version used, if any: any
* Python version, if any: any
* Operating System: any

### Is your feature request related to a problem? Please describe.

estimagic allows the criterion function to have two returns.

1. A criterion value or array which is passed to the optimizer.
2. A DataFrame containing additional information which can be visualized with the [comparison plot](https://estimagic.readthedocs.io/en/latest/visualization/comparison_plot_demo.html). #313 implements the same for the estimation via ML and allows to visualize the distribution of various contributions.

This issue tackles the second point. Passing data to the comparison plot would allow us to see how the model fit in terms of moments develops over the course of the optimization.

One could directly see which moments can be reproduced by the model and which ones not by comparing the value of the moments side-by-side or observing the difference between the two.

### Describe the solution you'd like

Something very similar #313 where the data adheres to the tidy data format.

The main point of discussion is what do we want to visualize. Some suggestions:

- The comparison plot data might have some key called ``kind`` which is either simulated or empirical.
- Do we want to visualize both, the simulated and empirical moment, or just the difference between the two.
- The moments can be grouped by which ``calc_moments`` produced them. In case, the functions are passed as dictionaries, we could re-use the dictionary keys. Otherwise, we have to use positions. This will allow us to group choice probabilities and wages, etc..

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=h1) Report
> Merging [#363](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=desc) into [master](https://codecov.io/gh/OpenSourceEconomics/respy/commit/4813ed80a07a4f6b304670ee6ba8226fab5cdafd&el=desc) will **decrease** coverage by `14.90%`.
> The diff coverage is `6.00%`.

[![Impacted file tree graph](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/graphs/tree.svg?width=650&height=150&src=pr&token=lefIFygbXJ)](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           master     #363       +/-   ##
===========================================
- Coverage   84.61%   69.70%   -14.91%     
===========================================
  Files          45       45               
  Lines        2957     3004       +47     
===========================================
- Hits         2502     2094      -408     
- Misses        455      910      +455     
```

| Flag | Coverage Δ | |
|---|---|---|
| #end_to_end_edge_case | `?` | |
| #end_to_end_precise | `?` | |
| #end_to_end_rest | `?` | |
| #integration_edge_case | `?` | |
| #integration_precise | `?` | |
| #integration_rest | `?` | |
| #unit_precise | `69.70% <6.00%> (-1.59%)` | :arrow_down: |
| #unit_rest | `?` | |

| [Impacted Files](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [respy/method\_of\_simulated\_moments.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvbWV0aG9kX29mX3NpbXVsYXRlZF9tb21lbnRzLnB5) | `12.39% <6.00%> (-67.34%)` | :arrow_down: |
| [respy/data.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvZGF0YS5weQ==) | `28.00% <0.00%> (-72.00%)` | :arrow_down: |
| [respy/tests/test\_simulate.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvdGVzdHMvdGVzdF9zaW11bGF0ZS5weQ==) | `30.76% <0.00%> (-69.24%)` | :arrow_down: |
| [respy/interpolate.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvaW50ZXJwb2xhdGUucHk=) | `24.63% <0.00%> (-68.12%)` | :arrow_down: |
| [respy/tests/test\_method\_of\_simulated\_moments.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvdGVzdHMvdGVzdF9tZXRob2Rfb2Zfc2ltdWxhdGVkX21vbWVudHMucHk=) | `34.04% <0.00%> (-65.96%)` | :arrow_down: |
| [respy/tests/test\_likelihood.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvdGVzdHMvdGVzdF9saWtlbGlob29kLnB5) | `34.48% <0.00%> (-65.52%)` | :arrow_down: |
| [respy/tests/test\_randomness.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvdGVzdHMvdGVzdF9yYW5kb21uZXNzLnB5) | `36.00% <0.00%> (-64.00%)` | :arrow_down: |
| [respy/tests/test\_integration.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvdGVzdHMvdGVzdF9pbnRlZ3JhdGlvbi5weQ==) | `36.84% <0.00%> (-63.16%)` | :arrow_down: |
| [respy/tests/test\_interpolate.py](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree#diff-cmVzcHkvdGVzdHMvdGVzdF9pbnRlcnBvbGF0ZS5weQ==) | `57.14% <0.00%> (-42.86%)` | :arrow_down: |
| ... and [13 more](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=footer). Last update [4813ed8...0418dee](https://codecov.io/gh/OpenSourceEconomics/respy/pull/363?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

amageh: Update on the second point. After feedback from @tobiasraabe, the comparison plot data now looks like this: 

2.  Comparison Plot Data: This feature is implemented through the argument `return_comparison_plot_data` which is boolean and returns a DataFrame in tidy data format. The returned DataFrame will have the following columns:

- **moment_column**: Contains the column names of the moment DataFrames/ Series names.
- **moment_index**: Contains the index of the moment DataFrames/Series (usually Periods).
- **value**: Contains the moment values.
- **moment_set**: Indicator for each set of moments, will use keys if `empirical_moments` is specified in a dict and will just be numbered otherwise.
- **kind**: Indicates whether moments are empirical or simulated.

I am still figuring out how to make this generally compatible with MultiIndex Series and DataFrames as moment containers. A two level Series would be very easy to allow for but introducing more index levels or columns on MultiIndex objects complicates things since those don't fit into the **moment_column**/**moment_index** structure anymore without joining the index or something.

@mo2561057 @tobiasraabe Let me know if you have any ideas/opinions about this, I will think about it some more. 

tobiasraabe: Fantastic job 👍 🎉 

Patch:
diff --git a/docs/how_to_guides/msm.ipynb b/docs/how_to_guides/msm.ipynb
index f448ac74..a95e86fb 100644
--- a/docs/how_to_guides/msm.ipynb
+++ b/docs/how_to_guides/msm.ipynb
@@ -46,6 +46,8 @@
     "* weighting_matrix (numpy.ndarray)\n",
     "* n_simulation_periods (int, default None)\n",
     "* return_scalar (bool, default True)\n",
+    "* return_simulated_moments (bool, default False)\n",
+    "* return_comparison_plot_data (bool, default False)\n",
     "\n",
     "\n",
     "`get_msm_func` returns a function where all arguments except *params* are held fixed. The returned function can then easily be passed on to an optimizer for estimation."
@@ -1109,6 +1111,24 @@
     "The *return_scalar* argument allows us to return the moment errors in vector form. `get_msm_func` will return the moment error vector if *return_scalar* is set to **False** and will return the value of the weighted square product of the moment errors if *return_scalar* is set to **True** which is also the default."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### The *return_simulated_moments* Argument\n",
+    "\n",
+    "If *return_simulated_moments* is set to **True**, the msm function will return the simulated in addition to the function value or moment error vector. This can be useful for visualization purposes and to compare simulated moments to empirical moments for a given parameterization. Its default value is **False**."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### The *return_comparison_plot_data* Argument\n",
+    "\n",
+    "Similarly to *return_simulated_moments*, the *return_comparison_plot_data* argument can be used to return the empirical and simulated moments with other function output. Instead of lists, the data will be saved to one pandas.DataFrame that is designed to compliment the visualization capabilities of [estimagic](https://github.com/OpenSourceEconomics/estimagic). It is **False** by default and acts as a substitute for *return_simulated_moments*, meaning only one can be set to **True** at the same time."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/docs/release_notes.rst b/docs/release_notes.rst
index 57e299da..bb1e9123 100644
--- a/docs/release_notes.rst
+++ b/docs/release_notes.rst
@@ -85,7 +85,8 @@ releases are available on `Anaconda.org
   pre-defined models.
 - :gh:`361` adds standard deviations of parameters for example models
   (:ghuser:`timmens`).
-
+- :gh:`363` enables msm function to return simulated moments or comparison plot data for 
+  use with `estimagic <https://github.com/OpenSourceEconomics/estimagic>`_ (:ghuser:`amageh`).
 
 *Releases prior to the second version were published on PyPI, but later deleted. You can
 still checkout the following releases using the corresponding tags in the repository.*
diff --git a/respy/method_of_simulated_moments.py b/respy/method_of_simulated_moments.py
index 127acb3d..b79a798b 100644
--- a/respy/method_of_simulated_moments.py
+++ b/respy/method_of_simulated_moments.py
@@ -34,8 +34,10 @@ def get_msm_func(
     weighting_matrix,
     n_simulation_periods=None,
     return_scalar=True,
+    return_simulated_moments=False,
+    return_comparison_plot_data=False,
 ):
-    """Get the msm function.
+    """Get the MSM function.
 
     Parameters
     ----------
@@ -68,8 +70,23 @@ def get_msm_func(
         number of periods for which decision rules are computed.
     return_scalar : bool, default True
         Indicates whether to return moment error vector (False) or weighted
-        square product of moment error vector (True).
-
+        square product of moment error vectors (True).
+    return_simulated_moments: bool, default False
+        Indicates whether simulated moments should be returned with other output.
+        If True will return simulated moments of the same type as empirical_moments.
+    return_comparison_plot_data: bool, default False
+        Indicator for whether a :class:`pandas.DataFrame` with empirical and simulated
+        moments for the visualization with estimagic should be returned. Data contains
+        the following columns:
+        - moment_column: Contains the column names of the moment DataFrames/Series
+        names.
+        - moment_index: Contains the index of the moment DataFrames/Series. MultiIndex
+        indices will be joined to one string.
+        - value: Contains moment values.
+        - moment_set: Indicator for each set of moments, will use keys if
+        empirical_moments are specified in a dict. Moments input as lists will
+        be numbered according to position.
+        - kind: Indicates whether moments are empirical or simulated.
     Returns
     -------
     msm_func: callable
@@ -78,6 +95,14 @@ def get_msm_func(
     """
     empirical_moments = copy.deepcopy(empirical_moments)
 
+    # Save keys of dictionary for comparison plot if applicable.
+    return_comparison_plot_data = [return_comparison_plot_data]
+    if isinstance(empirical_moments, dict):
+        moment_keys = sorted(empirical_moments)
+        return_comparison_plot_data.append(moment_keys)
+    else:
+        return_comparison_plot_data.append(None)
+
     simulate = get_simulate_func(
         params=params, options=options, n_simulation_periods=n_simulation_periods
     )
@@ -111,6 +136,11 @@ def get_msm_func(
             "the number of sets of empirical moments."
         )
 
+    if return_simulated_moments and return_comparison_plot_data[0]:
+        raise ValueError(
+            "Can only return either simulated moments or comparison plot data, not both."
+        )
+
     msm_func = functools.partial(
         msm,
         simulate=simulate,
@@ -119,6 +149,8 @@ def get_msm_func(
         empirical_moments=empirical_moments,
         weighting_matrix=weighting_matrix,
         return_scalar=return_scalar,
+        return_simulated_moments=return_simulated_moments,
+        return_comparison_plot_data=return_comparison_plot_data,
     )
 
     return msm_func
@@ -132,15 +164,17 @@ def msm(
     empirical_moments,
     weighting_matrix,
     return_scalar,
+    return_simulated_moments,
+    return_comparison_plot_data,
 ):
-    """Loss function for msm estimation.
+    """Loss function for MSM estimation.
 
     Parameters
     ----------
     params : pandas.DataFrame or pandas.Series
         Contains model parameters.
     simulate : callable
-        Function used to simulate data for msm estimation.
+        Function used to simulate data for MSM estimation.
     calc_moments : list
         List of function(s) used to calculate simulated moments. Must match length of
         empirical_moments i.e. calc_moments contains a moments function for each item in
@@ -159,11 +193,22 @@ def msm(
     return_scalar : bool
         Indicates whether to return moment error vector (False) or weighted square
         product of moment error vector (True).
+    return_simulated_moments: bool
+        Indicates whether simulated moments should be returned with other output.
+        If True will return simulated moments of the same type as empirical_moments.
+    return_comparison_plot_data: list
+        Will output moments in a tidy data format if True. Expects a list as input where
+        the first element is a boolean indicating whether to return the comparison plot
+        data. The second element in the list can be a list of keys used to identify sets
+        of moments which otherwise will be numbered.
 
     Returns
     -------
-    out : pandas.Series or float
-        Scalar or moment error vector depending on value of return_scalar.
+    out : pandas.Series or float or tuple
+        Scalar or moment error vector depending on value of return_scalar. Will be a
+        tuple containing simulated moments of same type as empirical_moments or a tidy
+        pandas.DataFrame if either return_simulated_moments or the first element in
+        return_comparison_plot_data is True.
 
     """
     empirical_moments = copy.deepcopy(empirical_moments)
@@ -193,6 +238,21 @@ def msm(
     else:
         out = moment_errors
 
+    if return_simulated_moments:
+        simulated_moments = _reconstruct_inputs(
+            simulated_moments, return_comparison_plot_data[1]
+        )
+        out = (out, simulated_moments)
+
+    elif return_comparison_plot_data[0]:
+        tidy_moments = _create_comparison_plot_data_msm(
+            empirical_moments, simulated_moments, return_comparison_plot_data[1]
+        )
+        out = (out, tidy_moments)
+
+    else:
+        pass
+
     return out
 
 
@@ -231,7 +291,7 @@ def get_diag_weighting_matrix(empirical_moments, weights=None):
         weights = _harmonize_input(weights)
 
         # Reindex weights to ensure they are assigned to the correct moments in
-        # the msm function.
+        # the MSM function.
         weights = [
             weight.reindex_like(emp_mom)
             for emp_mom, weight in zip(empirical_moments, weights)
@@ -275,7 +335,8 @@ def _harmonize_input(data):
       the dictionary entries.
 
     """
-    # Convert single DataFrames, Series or function into list containing one item.
+    # Convert single pandas.DataFrames, pandas.Series or function into list containing
+    # one item.
     if isinstance(data, (pd.DataFrame, pd.Series)) or callable(data):
         data = [data]
 
@@ -297,15 +358,18 @@ def _harmonize_input(data):
 
 def _flatten_index(data):
     """Flatten the index as a combination of the former index and the columns."""
+    data = copy.deepcopy(data)
     data_flat = []
     counter = itertools.count()
 
     for series_or_df in data:
         series_or_df.index = series_or_df.index.map(str)
-        # Unstack DataFrames and Series to add columns/Series name to index.
+        # Unstack pandas.DataFrames and pandas.Series to add
+        # columns/name to index.
         if isinstance(series_or_df, pd.DataFrame):
             df = series_or_df.rename(columns=str)
-        # Series without a name are named using a counter to avoid duplicate indexes.
+        # pandas.Series without a name are named using a counter to avoid duplicate
+        # indexes.
         elif isinstance(series_or_df, pd.Series) and series_or_df.name is None:
             df = series_or_df.to_frame(name=str(next(counter)))
         else:
@@ -317,3 +381,61 @@ def _flatten_index(data):
         data_flat.append(df)
 
     return pd.concat(data_flat)
+
+
+def _create_comparison_plot_data_msm(
+    empirical_moments, simulated_moments, moment_set_labels
+):
+    """Create pandas.DataFrame for estimagic's comparison plot."""
+    if moment_set_labels is None:
+        moment_set_labels = list(range(0, len(empirical_moments)))
+
+    tidy_empirical_moments = _create_tidy_data(empirical_moments, moment_set_labels)
+    tidy_simulated_moments = _create_tidy_data(simulated_moments, moment_set_labels)
+
+    tidy_simulated_moments["kind"] = "simulated"
+    tidy_empirical_moments["kind"] = "empirical"
+
+    return pd.concat(
+        [tidy_empirical_moments, tidy_simulated_moments], ignore_index=True
+    )
+
+
+def _create_tidy_data(data, moment_set_labels):
+    """Create tidy data from list of pandas.DataFrames."""
+    counter = itertools.count()
+    tidy_data = []
+    for series_or_df, label in zip(data, moment_set_labels):
+        # Join index levels for MultiIndex objects.
+        if isinstance(series_or_df.index, pd.MultiIndex):
+            series_or_df = series_or_df.rename(index=str)
+            series_or_df.index = series_or_df.index.to_flat_index().str.join("_")
+        # If moments are a pandas.Series, convert into pandas.DataFrame.
+        if isinstance(series_or_df, pd.Series):
+            # Unnamed pandas.Series receive a name based on a counter.
+            if series_or_df.name is None:
+                series_or_df = series_or_df.to_frame(name=next(counter))
+            else:
+                series_or_df = series_or_df.to_frame()
+
+        # Create pandas.DataFrame in tidy format.
+        tidy_df = series_or_df.unstack()
+        tidy_df.index.names = ("moment_column", "moment_index")
+        tidy_df.rename("value", inplace=True)
+        tidy_df = tidy_df.reset_index()
+        tidy_df["moment_set"] = label
+        tidy_data.append(tidy_df)
+
+    return pd.concat(tidy_data, ignore_index=True)
+
+
+def _reconstruct_inputs(inputs, dict_keys=None):
+    """Reconstruct inputs from lists back to a dictionary or single object."""
+    if dict_keys is not None:
+        output = dict(zip(dict_keys, inputs))
+    elif len(inputs) == 1:
+        output = inputs[0]
+    else:
+        output = inputs
+
+    return output


Test patch:
diff --git a/respy/tests/test_method_of_simulated_moments.py b/respy/tests/test_method_of_simulated_moments.py
index 427df713..c6d8b00f 100644
--- a/respy/tests/test_method_of_simulated_moments.py
+++ b/respy/tests/test_method_of_simulated_moments.py
@@ -1,4 +1,5 @@
 """Test the msm interface of respy."""
+import pandas as pd
 import pytest
 
 from respy.interface import get_example_model
@@ -93,6 +94,35 @@ def test_randomness_msm(model_or_seed):
     assert msm(params) == 0
 
 
+@pytest.mark.integration
+def test_return_simulated_moments_for_msm(inputs):
+    """Return_simulated_moments."""
+    msm = get_msm_func(*inputs, return_simulated_moments=True)
+    fval, simulated_moments = msm(inputs[0])
+
+    assert isinstance(fval, float)
+    assert isinstance(simulated_moments, (dict, list, pd.DataFrame, pd.Series))
+
+
+@pytest.mark.integration
+def test_return_comparison_plot_data_for_msm(inputs):
+    """Return_comparison_plot_data."""
+    msm = get_msm_func(*inputs, return_scalar=False, return_comparison_plot_data=True)
+    moment_errors, df = msm(inputs[0])
+
+    assert isinstance(moment_errors, pd.Series)
+    assert isinstance(df, pd.DataFrame)
+
+
+@pytest.mark.integration
+def test_multiple_returns_msm(inputs):
+    """Raise error if moments and comparison plot data is requested."""
+    with pytest.raises(ValueError, match="Can only return either"):
+        get_msm_func(
+            *inputs, return_simulated_moments=True, return_comparison_plot_data=True
+        )
+
+
 def _calc_choice_freq(df):
     return df.groupby("Period").Choice.value_counts(normalize=True).unstack()