add stokes problem

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]),
+    }

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)
+
+