.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/minipyro.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_minipyro.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_minipyro.py:


Example: Mini Pyro
==================

.. GENERATED FROM PYTHON SOURCE LINES 9-83

.. code-block:: Python


    import argparse

    import torch
    from pyroapi import distributions as dist
    from pyroapi import infer, optim, pyro, pyro_backend
    from torch.distributions import constraints

    import funsor
    from funsor.montecarlo import MonteCarlo


    def main(args):
        funsor.set_backend("torch")

        # Define a basic model with a single Normal latent random variable `loc`
        # and a batch of Normally distributed observations.
        def model(data):
            loc = pyro.sample("loc", dist.Normal(0.0, 1.0))
            with pyro.plate("data", len(data), dim=-1):
                pyro.sample("obs", dist.Normal(loc, 1.0), obs=data)

        # Define a guide (i.e. variational distribution) with a Normal
        # distribution over the latent random variable `loc`.
        def guide(data):
            guide_loc = pyro.param("guide_loc", torch.tensor(0.0))
            guide_scale = pyro.param(
                "guide_scale", torch.tensor(1.0), constraint=constraints.positive
            )
            pyro.sample("loc", dist.Normal(guide_loc, guide_scale))

        # Generate some data.
        torch.manual_seed(0)
        data = torch.randn(100) + 3.0

        # Because the API in minipyro matches that of Pyro proper,
        # training code works with generic Pyro implementations.
        with pyro_backend(args.backend), MonteCarlo():
            # Construct an SVI object so we can do variational inference on our
            # model/guide pair.
            Elbo = infer.JitTrace_ELBO if args.jit else infer.Trace_ELBO
            elbo = Elbo()
            adam = optim.Adam({"lr": args.learning_rate})
            svi = infer.SVI(model, guide, adam, elbo)

            # Basic training loop
            pyro.get_param_store().clear()
            for step in range(args.num_steps):
                loss = svi.step(data)
                if args.verbose and step % 100 == 0:
                    print(f"step {step} loss = {loss}")

            # Report the final values of the variational parameters
            # in the guide after training.
            if args.verbose:
                for name in pyro.get_param_store():
                    value = pyro.param(name).data
                    print(f"{name} = {value.detach().cpu().numpy()}")

            # For this simple (conjugate) model we know the exact posterior. In
            # particular we know that the variational distribution should be
            # centered near 3.0. So let's check this explicitly.
            assert (pyro.param("guide_loc") - 3.0).abs() < 0.1


    if __name__ == "__main__":
        parser = argparse.ArgumentParser(description="Minipyro demo")
        parser.add_argument("-b", "--backend", default="funsor")
        parser.add_argument("-n", "--num-steps", default=1001, type=int)
        parser.add_argument("-lr", "--learning-rate", default=0.02, type=float)
        parser.add_argument("--jit", action="store_true")
        parser.add_argument("-v", "--verbose", action="store_true")
        args = parser.parse_args()
        main(args)


.. _sphx_glr_download_examples_minipyro.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: minipyro.ipynb <minipyro.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: minipyro.py <minipyro.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: minipyro.zip <minipyro.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_