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

.. only:: html

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

        :ref:`Go to the end <sphx_glr_download_examples_ex16_deeplearning-torch.py>`
        to download the full example code.

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

.. _sphx_glr_examples_ex16_deeplearning-torch.py:


Train a deep neural network
---------------------------

.. topic:: Train a PyTorch model by simulation results.

   * compression of a hyperelastic axisymmetric cylinder

   * evaluate displacements located at mesh-points

   * train a neural network on the displacement data

   * obtain displacements from the PyTorch model and plot the log. strains

First, an axisymmetric model is created. The displacements are saved after each
completed substep. Only very few substeps are used to run the simulation.

.. admonition:: This example requires external packages.
   :class: hint

   .. code-block::

      pip install torch

.. GENERATED FROM PYTHON SOURCE LINES 25-46

.. code-block:: Python


    import numpy as np

    import felupe as fem

    mesh = fem.Rectangle(a=(0, 1), b=(1, 3), n=(11, 31))
    region = fem.RegionQuad(mesh)
    field = fem.FieldContainer([fem.FieldAxisymmetric(region, dim=2)])
    boundaries = fem.dof.uniaxial(field, clamped=True, sym=False, return_loadcase=False)
    solid = fem.SolidBody(fem.NeoHookeCompressible(mu=1, lmbda=10), field)
    move = fem.math.linsteps([0, -0.2], num=3)
    step = fem.Step(items=[solid], ramp={boundaries["move"]: move}, boundaries=boundaries)
    displacements = np.zeros((len(move), *field[0].values.shape))


    def save(i, j, res):
        displacements[j] = res.x[0].values


    job = fem.Job(steps=[step], callback=save).evaluate(tol=1e-1)








.. GENERATED FROM PYTHON SOURCE LINES 48-50

A PyTorch model is trained on the simulation results. For simplicity, testing is
skipped and the data is not splitted in batches.

.. GENERATED FROM PYTHON SOURCE LINES 50-101

.. code-block:: Python

    import torch
    import torch.nn as nn
    from tqdm import tqdm


    class NeuralNetwork(nn.Module):
        def __init__(self, in_features, out_features):
            super().__init__()
            self.flatten = nn.Flatten()
            self.linear_relu_stack = nn.Sequential(
                nn.Linear(in_features, 128),
                nn.ReLU(),
                nn.Linear(128, 64),
                nn.ReLU(),
                nn.Linear(64, out_features),
            )

        def forward(self, x):
            x = self.flatten(x)
            logits = self.linear_relu_stack(x)
            return logits.reshape(len(x), mesh.npoints, mesh.dim)


    xv, yv = move.reshape(-1, 1), displacements
    x = torch.tensor(xv, dtype=torch.float32)
    y = torch.tensor(yv, dtype=torch.float32)

    model = NeuralNetwork(in_features=xv[0].size, out_features=yv[0].size)

    loss_fn = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters())


    def train(x, y, model, loss_fn, optimizer, verbose=False):
        model.train()
        pred = model(x)
        loss = loss_fn(pred, y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        if verbose:
            print(f"Train Loss (Final): {loss.item():.2e}")


    def run_epochs(n=100):
        for t in tqdm(range(n - 1)):
            train(x, y, model, loss_fn, optimizer)
        train(x, y, model, loss_fn, optimizer, verbose=True)
        model.eval()









.. GENERATED FROM PYTHON SOURCE LINES 102-105

Max. principal values of the logarithmic strain tensors are evaluated and plotted
based on the PyTorch model displacement results. After 50 epoch runs, the result is
quite bad with a lot of unwanted artefacts (noise).

.. GENERATED FROM PYTHON SOURCE LINES 105-110

.. code-block:: Python

    run_epochs(n=50)
    field_2 = field.copy()
    field_2[0].values[:] = model(torch.Tensor([[-0.17]])).detach().numpy()[0]
    field_2.plot("Principal Values of Logarithmic Strain").show()








.. tab-set::



   .. tab-item:: Static Scene



            
     .. image-sg:: /examples/images/sphx_glr_ex16_deeplearning-torch_001.png
        :alt: ex16 deeplearning torch
        :srcset: /examples/images/sphx_glr_ex16_deeplearning-torch_001.png
        :class: sphx-glr-single-img
     


   .. tab-item:: Interactive Scene



       .. offlineviewer:: /home/docs/checkouts/readthedocs.org/user_builds/felupe/checkouts/stable/docs/examples/images/sphx_glr_ex16_deeplearning-torch_001.vtksz



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

 .. code-block:: none


      0%|          | 0/49 [00:00<?, ?it/s]
     37%|███▋      | 18/49 [00:00<00:00, 178.38it/s]
    100%|██████████| 49/49 [00:00<00:00, 372.65it/s]
    Train Loss (Final): 1.23e-03




.. GENERATED FROM PYTHON SOURCE LINES 111-112

After 500 more epoch runs, the result is much more realistic.

.. GENERATED FROM PYTHON SOURCE LINES 112-115

.. code-block:: Python

    run_epochs(n=500)
    field_2[0].values[:] = model(torch.Tensor([[-0.17]])).detach().numpy()[0]
    field_2.plot("Principal Values of Logarithmic Strain").show()







.. tab-set::



   .. tab-item:: Static Scene



            
     .. image-sg:: /examples/images/sphx_glr_ex16_deeplearning-torch_002.png
        :alt: ex16 deeplearning torch
        :srcset: /examples/images/sphx_glr_ex16_deeplearning-torch_002.png
        :class: sphx-glr-single-img
     


   .. tab-item:: Interactive Scene



       .. offlineviewer:: /home/docs/checkouts/readthedocs.org/user_builds/felupe/checkouts/stable/docs/examples/images/sphx_glr_ex16_deeplearning-torch_002.vtksz



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

 .. code-block:: none


      0%|          | 0/499 [00:00<?, ?it/s]
     19%|█▉        | 95/499 [00:00<00:00, 947.92it/s]
     38%|███▊      | 191/499 [00:00<00:00, 951.58it/s]
     58%|█████▊    | 287/499 [00:00<00:00, 945.75it/s]
     77%|███████▋  | 382/499 [00:00<00:00, 946.90it/s]
     96%|█████████▌| 478/499 [00:00<00:00, 949.66it/s]
    100%|██████████| 499/499 [00:00<00:00, 948.22it/s]
    Train Loss (Final): 7.48e-07





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

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


.. _sphx_glr_download_examples_ex16_deeplearning-torch.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: ex16_deeplearning-torch.ipynb <ex16_deeplearning-torch.ipynb>`

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

      :download:`Download Python source code: ex16_deeplearning-torch.py <ex16_deeplearning-torch.py>`

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

      :download:`Download zipped: ex16_deeplearning-torch.zip <ex16_deeplearning-torch.zip>`


.. only:: html

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

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