diff --git a/examples/problems/stokes.py b/examples/problems/stokes.py new file mode 100644 index 0000000..9874f29 --- /dev/null +++ b/examples/problems/stokes.py @@ -0,0 +1,45 @@ +import numpy as np +import torch + +from pina.problem import SpatialProblem +from pina.operators import nabla, grad, div +from pina import Condition, Span, LabelTensor + + +class Stokes(SpatialProblem): + + spatial_variables = ['x', 'y'] + output_variables = ['ux', 'uy', 'p'] + domain = Span({'x': [-2, 2], 'y': [-1, 1]}) + + def momentum(input_, output_): + #print(nabla(output_['ux', 'uy'], input_)) + #print(grad(output_['p'], input_)) + nabla_ = LabelTensor.hstack([ + LabelTensor(nabla(output_['ux'], input_), ['x']), + LabelTensor(nabla(output_['uy'], input_), ['y'])]) + #return LabelTensor(nabla_.tensor + grad(output_['p'], input_).tensor, ['x', 'y']) + return nabla_.tensor + grad(output_['p'], input_).tensor + + def continuity(input_, output_): + return div(output_['ux', 'uy'], input_) + + def inlet(input_, output_): + value = 2.0 + return output_['ux'] - value + + def outlet(input_, output_): + value = 0.0 + return output_['p'] - value + + def wall(input_, output_): + value = 0.0 + return output_['ux', 'uy'].tensor - value + + conditions = { + 'gamma_top': Condition(Span({'x': [-2, 2], 'y': 1}), wall), + 'gamma_bot': Condition(Span({'x': [-2, 2], 'y': -1}), wall), + 'gamma_out': Condition(Span({'x': 2, 'y': [-1, 1]}), outlet), + 'gamma_in': Condition(Span({'x': -2, 'y': [-1, 1]}), inlet), + 'D': Condition(Span({'x': [-2, 2], 'y': [-1, 1]}), [momentum, continuity]), + } diff --git a/examples/run_stokes.py b/examples/run_stokes.py new file mode 100644 index 0000000..e6438cd --- /dev/null +++ b/examples/run_stokes.py @@ -0,0 +1,54 @@ +import argparse +import sys +import numpy as np +import torch +from torch.nn import ReLU, Tanh, Softplus + +from pina import PINN, LabelTensor, Plotter +from pina.model import FeedForward +from pina.adaptive_functions import AdaptiveSin, AdaptiveCos, AdaptiveTanh +from problems.stokes import Stokes + +if __name__ == "__main__": + + parser = argparse.ArgumentParser(description="Run PINA") + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("-s", "-save", action="store_true") + group.add_argument("-l", "-load", action="store_true") + parser.add_argument("id_run", help="number of run", type=int) + args = parser.parse_args() + + + stokes_problem = Stokes() + model = FeedForward( + layers=[40, 20, 20, 10], + output_variables=stokes_problem.output_variables, + input_variables=stokes_problem.input_variables, + func=Softplus, + ) + + pinn = PINN( + stokes_problem, + model, + lr=0.006, + error_norm='mse', + regularizer=1e-8, + lr_accelerate=None) + + if args.s: + + #pinn.span_pts(200, 'grid', ['gamma_out']) + pinn.span_pts(200, 'grid', ['gamma_top', 'gamma_bot', 'gamma_in', 'gamma_out']) + pinn.span_pts(2000, 'random', ['D']) + #plotter = Plotter() + #plotter.plot_samples(pinn) + pinn.train(10000, 100) + pinn.save_state('pina.stokes') + + else: + pinn.load_state('pina.stokes') + plotter = Plotter() + plotter.plot_samples(pinn) + plotter.plot(pinn) + +