<jupyter_start><jupyter_text>Deep Ritz: PoissonSolving the poisson equation in 2D using the Deep Ritz Method.<jupyter_code>import torch
import torchphysics as tp<jupyter_output><empty_output><jupyter_text>First, we create the spaces for our problem. These define the variable names which will be used in the remaining part of this code.In this example, x is the space variable and u is the variable for the (1D-)solution. For more complex problems, spaces could also be built as cartesian products using the * operator<jupyter_code>X = tp.spaces.R2('x')
U = tp.spaces.R1('u')<jupyter_output><empty_output><jupyter_text>As a next step, we build the domain of the problem. There are multiple options to build multi-dimensional domains - in this case, we simply create a circle and cut a square in its center.<jupyter_code>A = tp.domains.Circle(X, center=[0, 0], radius=1.0)
B = tp.domains.Parallelogram(X, [-0.25, -0.25], [-0.25, 0.25], [0.25, -0.25])

# cutting operation:
D = A-B<jupyter_output><empty_output><jupyter_text>Before we visualize the created domain, we create Sampler objects which are iterators that sample points from the domain during the optimization task. There are again various options to sample from the domains, an easy way would be to sample uniformly distributed random points.<jupyter_code>inner_sampler = tp.samplers.RandomUniformSampler(D, density=1e4)
boundary_sampler = tp.samplers.RandomUniformSampler(D.boundary, density=1e4)<jupyter_output><empty_output><jupyter_text>We visualize the domain through the points created by the samplers using matplotlib:<jupyter_code># to plot the points, really memory/computional expensive to plot!
"""
inner_points = next(inner_sampler)
boundary_points = next(boundary_sampler)

from matplotlib import pyplot as plt

plt.scatter(inner_points.as_tensor[:,0],
            inner_points.as_tensor[:,1])

plt.scatter(boundary_points.as_tensor[:,0],
            boundary_points.as_tensor[:,1])
"""<jupyter_output><empty_output><jupyter_text>In the next step we define the NN-model we want to fit to the PDE and a Condition which aims to minimze the variational part of the deep ritz method.<jupyter_code>model = tp.models.FCN(input_space=X, output_space=U, hidden=(50,50,50))

def poisson_integrand(u, x):
    a = 0.5*(torch.sum(tp.utils.grad(u, x)**2, dim=1))
    return a

pde_condition = tp.conditions.DeepRitzCondition(module=model,
                                                sampler=inner_sampler,
                                                integrand_fn=poisson_integrand,
                                                weight=5)<jupyter_output><empty_output><jupyter_text>Additionally, we add a boundary condition at the boundary of the domain:<jupyter_code>def f(x):
    return torch.cos(x[:, :1])*torch.cos(x[:, 1:])

def boundary_residual(u, f):
    return (u - f)**2

boundary_condition = tp.conditions.DeepRitzCondition(module=model,
                                                     sampler=boundary_sampler,
                                                     integrand_fn=boundary_residual,
                                                     data_functions={'f': f})<jupyter_output><empty_output><jupyter_text>Finally, we optimize the conditions using a pytorch-lightning.LightningModule Solver and running the training. In the Solver, the training and validation conditions, as well as all optimizer options can be specified.<jupyter_code>solver = tp.solver.Solver([pde_condition,
                           boundary_condition])

import pytorch_lightning as pl
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

trainer = pl.Trainer(gpus=None, # or None for CPU
                     max_steps=2000,
                     logger=False,
                     benchmark=True,
                     checkpoint_callback=False)
trainer.fit(solver)<jupyter_output>GPU available: True, used: False
TPU available: False, using: 0 TPU cores
/home/tomfre/miniconda3/envs/bosch/lib/python3.9/site-packages/pytorch_lightning/utilities/distributed.py:69: UserWarning: GPU available but not used. Set the gpus flag in your trainer `Trainer(gpus=1)` or script `--gpus=1`.
  warnings.warn(*args, **kwargs)

  | Name             | Type       | Params
------------------------------------------------
0 | train_conditions | ModuleList | 5.3 K 
1 | val_conditions   | ModuleList | 0     
------------------------------------------------
5.3 K     Trainable params
0         Non-trainable params
5.3 K     Total params
0.021     Total estimated model params size (MB)<jupyter_text>Finally, we plot the obtained solution:<jupyter_code>tp.utils.plot(model, lambda u: u[:, 0], tp.samplers.PlotSampler(D, density=100))
# the handed lambda function could also used to plot the laplacian
tp.utils.plot(model, lambda u, x: tp.utils.laplacian(u, x), tp.samplers.PlotSampler(D, density=100))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11102007374027038, 'fraction_numerical': 0.016796394920114706, 'mean_word_length': 3.180650684931507, 'pattern_counts': {'":': 0, '<': 36, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '137533', 'license': 'Apache-2.0', 'n_tokens_mistral': 1500, 'n_tokens_neox': 1454, 'n_words': 422, 'path': 'examples/deepritz/poisson-equation.ipynb', 'repo_name': 'uwe-iben/torchphysics', 'size': '4882'}