Add plots to docstrings of experiment classes

[drbenvincent]

NOTE: We are just about to have a major refactor of the codebase, so this should be done after #381 is merged.

It might be nice to make the api docs (auto build from the docstrings) more engaging by adding some plots. As far as I can tell, it is possible to take the code that we have in the doctests right now and also generate a plot that is rendered in the docs.

This could be relevant for the experiment classes in particular. I don't see any value in adding to the model classes for example.

Here's an example of current docs:

If we add the plot method into that code snippet, then we can get that plot to render in the built docs.

I believe (though have not verified) that we can do this as follows:

.. plot::
    :context: close-figs
    :include-source: True
    :format: doctest

    >>> import causalpy as cp
    >>> df = cp.load_data("did")
    >>> seed = 42
    >>> result = cp.pymc_experiments.DifferenceInDifferences(
    ...     df,
    ...     formula="y ~ 1 + group*post_treatment",
    ...     time_variable_name="t",
    ...     group_variable_name="group",
    ...     model=cp.pymc_models.LinearRegression(
    ...         sample_kwargs={
    ...             "target_accept": 0.95,
    ...             "random_seed": seed,
    ...             "progressbar": False,
    ...         }
    ...     )
    ...  )
    >>> result.plot()

If that works as I imagine, then this will:

• work as a doctest
• render the code in the docs
• include the rendered plot in the docs.

I have not verified that this works, but I imagine there is a way!

Notes

1. We might need to call plt.show() at the end.
2. At this point I have no idea if the doctest work would be taken over by sphinx, because there is an extension sphinx.ext.doctest. Would need to look into this.
3. We'd also need to add sphinx.ext.plot_directive to the list of extensions in conf.py.

[harsha-mangena]

Hey @drbenvincent, Can I pick this issue?

[drbenvincent]

Hi @harsha-mangena. Sounds good - but we're just about to merge a PR which is a pretty major code refactor. So we need to wait until #381 is merged first or any work done will be messed up. Can you check back toward the end of next week? It should (hopefully) be merged then.

[harsha-mangena]

@drbenvincent, No problem. I will check this issue next week.