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

.. only:: html

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

        Click :ref:`here <sphx_glr_download_pgo_pgo_tutorial.py>`
        to download the full example code

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

.. _sphx_glr_pgo_pgo_tutorial.py:


Pose Graph Optimization Tutorial
================================

.. GENERATED FROM PYTHON SOURCE LINES 6-22

.. code-block:: default



    import os
    import torch
    import argparse
    import pypose as pp
    from torch import nn
    from pgo_dataset_tutorial import G2OPGO
    import matplotlib.pyplot as plt
    import pypose.optim.solver as ppos
    import pypose.optim.kernel as ppok
    import pypose.optim.corrector as ppoc
    import pypose.optim.strategy as ppost
    from pypose.optim.scheduler import StopOnPlateau









.. GENERATED FROM PYTHON SOURCE LINES 23-25

Define Pose Graph
------------------

.. GENERATED FROM PYTHON SOURCE LINES 25-70

.. code-block:: default


    class PoseGraph(nn.Module):

        def __init__(self, nodes):
            super().__init__()
            self.nodes = pp.Parameter(nodes)

        def forward(self, edges, poses):
            node1 = self.nodes[edges[..., 0]]
            node2 = self.nodes[edges[..., 1]]
            error = poses.Inv() @ node1.Inv() @ node2
            return error.Log().tensor()


    @torch.no_grad()
    def plot_and_save(points, pngname, title='', axlim=None):
        points = points.detach().cpu().numpy()
        plt.figure(figsize=(7, 7))
        ax = plt.axes(projection='3d')
        ax.plot3D(points[:,0], points[:,1], points[:,2], 'b')
        plt.title(title)
        if axlim is not None:
            ax.set_xlim(axlim[0])
            ax.set_ylim(axlim[1])
            ax.set_zlim(axlim[2])
        plt.savefig(pngname)
        print('Saving to', pngname)
        return ax.get_xlim(), ax.get_ylim(), ax.get_zlim()


    parser = argparse.ArgumentParser(description='Pose Graph Optimization')
    parser.add_argument("--device", type=str, default='cuda:0', help="cuda or cpu")
    parser.add_argument("--radius", type=float, default=1e4, help="trust region radius")
    parser.add_argument("--save", type=str, default='../dataset/pgo/save/', help="files location to save")
    parser.add_argument("--dataroot", type=str, default='../dataset/pgo', help="dataset location")
    parser.add_argument("--dataname", type=str, default='parking-garage.g2o', help="dataset name")
    parser.add_argument('--no-vectorize', dest='vectorize', action='store_false', help="to save memory")
    parser.add_argument('--vectorize', action='store_true', help='to accelerate computation')
    parser.set_defaults(vectorize=False)
    args = parser.parse_args(); print(args)
    os.makedirs(os.path.join(args.save), exist_ok=True)

    data = G2OPGO(args.dataroot, args.dataname, device=args.device)
    edges, poses, infos = data.edges, data.poses, data.infos





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Namespace(device='cuda:0', radius=10000.0, save='../dataset/pgo/save/', dataroot='../dataset/pgo', dataname='parking-garage.g2o', vectorize=False)




.. GENERATED FROM PYTHON SOURCE LINES 71-73

Define Optimizer
-------------------------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 73-83

.. code-block:: default


    graph = PoseGraph(data.nodes).to(args.device)
    solver = ppos.Cholesky()
    strategy = ppost.TrustRegion(radius=args.radius)
    optimizer = pp.optim.LM(graph, solver=solver, strategy=strategy, min=1e-6, vectorize=args.vectorize)
    scheduler = StopOnPlateau(optimizer, steps=10, patience=3, decreasing=1e-3, verbose=True)

    pngname = os.path.join(args.save, args.dataname+'.png')
    axlim = plot_and_save(graph.nodes.translation(), pngname, args.dataname)




.. image-sg:: /pgo/images/sphx_glr_pgo_tutorial_001.png
   :alt: parking-garage.g2o
   :srcset: /pgo/images/sphx_glr_pgo_tutorial_001.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Saving to ../dataset/pgo/save/parking-garage.g2o.png




.. GENERATED FROM PYTHON SOURCE LINES 84-87

the 1st implementation: for customization and easy to extend
-------------------------------------------------------------
commented out because too time consumming

.. GENERATED FROM PYTHON SOURCE LINES 87-97

.. code-block:: default


    # while scheduler.continual:
    #     loss = optimizer.step(input=(edges, poses), weight=infos)
    #     scheduler.step(loss)

    #     name = os.path.join(args.save, args.dataname + '_' + str(scheduler.steps))
    #     title = 'PyPose PGO at the %d step(s) with loss %7f'%(scheduler.steps, loss.item())
    #     plot_and_save(graph.nodes.translation(), name+'.png', title, axlim=axlim)
    #     torch.save(graph.state_dict(), name+'.pt')








.. GENERATED FROM PYTHON SOURCE LINES 98-100

The 2nd implementation: equivalent to the 1st one, but more compact
--------------------------------------------------------------------

.. GENERATED FROM PYTHON SOURCE LINES 100-102

.. code-block:: default


    # scheduler.optimize(input=(edges, poses), weight=infos)








.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.223 seconds)


.. _sphx_glr_download_pgo_pgo_tutorial.py:

.. only:: html

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


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

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

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

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


.. only:: html

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

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