<jupyter_start><jupyter_text>Solve a heat equation on a domain defined by an .stl fileThis example gives the code for solving the heat equation on a cornerplate.<jupyter_code>import os
import torch
import time
from torch._C import device
import pytorch_lightning as pl
import numpy as np

from torchphysics.problem.domain.domain3D import Polygon3D
from torchphysics.problem.domain.domain1D import Interval
from torchphysics.problem.condition import (DirichletCondition,
                                              DiffEqCondition,
                                              DataCondition)
from torchphysics.problem import Variable
from torchphysics.utils import laplacian, grad
from torchphysics import PINNModule, Setting
from torchphysics.models.fcn import SimpleFCN
os.environ["CUDA_VISIBLE_DEVICES"] = "0"<jupyter_output>/home/nick7/anaconda3/envs/bosch-pde/lib/python3.9/site-packages/pytorch_lightning/metrics/__init__.py:43: LightningDeprecationWarning: `pytorch_lightning.metrics.*` module has been renamed to `torchmetrics.*` and split off to its own package (https://github.com/PyTorchLightning/metrics) since v1.3 and will be removed in v1.5
  rank_zero_deprecation(<jupyter_text>We load the domain as a 3D volume from an .stl file. Afterwards we project it onto the xy-plane.<jupyter_code>from torchphysics.problem.domain.domain2D import Circle
from torchphysics.problem.domain.domain_operations import Cut

#p = Polygon3D(file_name='spalt_ex2_01.stl') #
o = Circle([0,0],10.0)
i = Circle([0,-2.5],5.0)
p = Cut(o,i)
#p = Polygon3D(file_name='cornerplate.stl')
#p = p.project_on_plane(plane_origin=[0,0,0], plane_normal=[0,0,1])
#p = p.project_on_plane(plane_origin=[0,0,0], plane_normal=[0,0,1])

import matplotlib.pyplot as plt
b = p.sample_boundary(800, type='grid')
print(p.surface)
plt.scatter(b[:, 0], b[:, 1])<jupyter_output>94.24777960769379<jupyter_text>The computed 2D polygon becomes the domain of x. We use D as an additional input to the network, just as we do in `heat_equation_D_variable.ipynb`. The behaviour of the code is also pretty similar to the other example.<jupyter_code># Create variables
x = Variable(name='x', order=2, domain=p,
             train_conditions={}, val_conditions={})
t = Variable(name='t', order=1, domain=Interval(0,1),
             train_conditions={}, val_conditions={})
D = Variable(name='D',
             order=0,
             domain=Interval(low_bound=10,
                             up_bound=20),
             train_conditions={},
             val_conditions={})
# Create boundary condition for x
norm = torch.nn.MSELoss()
def x_dirichlet_fun(input):
    return np.zeros_like(input['t'])
x.add_train_condition(DirichletCondition(dirichlet_fun=x_dirichlet_fun,
                                         name='dirichlet',
                                         norm=norm,
                                         dataset_size=800,
                                         data_plot_variables=('x', 't')))<jupyter_output><empty_output><jupyter_text>The inital values are chosen such that they are non-zero on a third of the domain, and zero on the remaining parts. They are therefore non-smooth.<jupyter_code># Create boundary condition for t, i.e. initial values
def t_dirichlet_fun(input):
    temp = np.zeros((input['x'][:, :1].shape[0], 1), dtype=np.float32)
    for i in range(temp.shape[0]):
        temp[i][0] = 1*(np.exp(-0.1*np.abs(input['x'][i,0]-0))
                        *np.exp(-0.1*np.abs(input['x'][i,1]-0)))
    return temp
t.add_train_condition(DirichletCondition(dirichlet_fun=t_dirichlet_fun,
                                         name='dirichlet',
                                         norm=norm,
                                         dataset_size=2500,
                                         boundary_sampling_strategy='lower_bound_only',
                                         data_plot_variables=('x', 't')))
# Create pde condition 
def pde(u, input):
    return grad(u, input['t']) - input['D']*laplacian(u, input['x'])
train_cond = DiffEqCondition(pde=pde,
                             norm=norm,
                             dataset_size=2000,
                             data_plot_variables=('x', 't'))
# Put everything together
from torchphysics.utils.plot import Plotter

setup = Setting(variables=(x, t, D),
                train_conditions={'pde': train_cond},
                val_conditions={},
                n_iterations=100)

plotter = Plotter(plot_variables=setup.variables['x'],
                  points=400,
                  dic_for_other_variables={'t': 1.0, 'D': 15.0},
                  all_variables=setup.variables,
                  log_interval=10)

model = SimpleFCN(input_dim=4, width=50)

solver = PINNModule(model=model,
                    optimizer=torch.optim.Adam,
                    lr=1e-3)
# Start training
zeit1 = time.time()
trainer = pl.Trainer(gpus='-1' if torch.cuda.is_available() else None,
                     max_epochs=20)
trainer.fit(solver, setup)
zeit2 = time.time()
print('Zeit = %.2f', zeit2-zeit1)
# Plot the solution
D = 10
alpha = 30 #30
beta  = 20 #220
from torchphysics.utils.plot import _plot
fig = _plot(model=model, plot_variables=x, points=600, angle=[alpha, beta], 
            dic_for_other_variables={'t' : 0, 'D' : D})
fig = _plot(model=model, plot_variables=x, points=600, angle=[alpha, beta], 
            dic_for_other_variables={'t' : 1, 'D' : D})

# Create a animation of the solution
from torchphysics.utils.animation import animation
fig, ani = animation(model=model, plot_variables=x, domain_points=200, angle=[30, 220], 
                animation_variable=t, frame_number=100, dic_for_other_variables={'D' : 10})
ani.save('heat_equation_D=10.gif')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1148766075773375, 'fraction_numerical': 0.027980535279805353, 'mean_word_length': 2.808735936465917, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '213770', 'license': 'Apache-2.0', 'n_tokens_mistral': 1866, 'n_tokens_neox': 1756, 'n_words': 407, 'path': 'experiments/geometry/heat_equation_bearing.ipynb', 'repo_name': 'uwe-iben/torchphysics', 'size': '5754'}