├── .gitignore
├── README.md
├── fig
├── airfoil.gif
├── flag.gif
├── gaussian_circle_square.gif
├── gaussian_square.gif
├── gaussian_square_dynamic.gif
├── possion_circle.png
├── possion_l.png
├── possion_square.png
├── t-shirt.gif
└── tearing.gif
├── notebook
├── ARCSim
│ └── ArcSim.ipynb
├── FEniCS
│ ├── Gaussian.ipynb
│ └── Poisson.ipynb
└── SU2
│ └── SU2 to DGL Tutorial.ipynb
└── src
├── __init__.py
├── gaussian.py
├── possion.py
└── utils
├── __init__.py
├── arcsim.py
├── bc.py
├── fenics.py
├── plot.py
├── su2.py
└── to_dgl.py
/.gitignore:
--------------------------------------------------------------------------------
1 | # Exclude generated data
2 | data/
3 |
4 | # Exclude extra files
5 | __pycache__
6 | .ipynb_checkpoints
7 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # 📦 PDEs Dataset Generator
2 |
3 | A tool for generating datasets for Partial Differential Equations (PDEs) from ARCSim, FEniCS and SU2.
4 |
5 | ## 💡 Introduction
6 |
7 | Research on PDEs needs ground truth datasets. Several tools have been developed and among all, the open-source FEniCS is one of the most famous ones. Our goal is to provide tools for more easily creating such datasets and save them in the DGL graph format.
8 |
9 | Features:
10 |
11 | - Compressed scripts containing full generating process, including creating domain, generating mesh, creating boundary constrains, solve function, transfer mesh and result to dgl graph and saving
12 | - Results are transferred to DGL graph, which is convenient to use in graph models; we also provide numpy tools
13 | - Detailed tutorials with notebooks to show how our scripts work and knowledge about using ARCSim, FEniCS and SU2
14 | - Useful mini tools, including mesh to DGL transferring and modified dgl graph plotting
15 |
16 | ## 🔧 Environment
17 |
18 | ```Python
19 | Python Version: 3.6 or later
20 | Python Packages: jupyterlab, fenics, dgl, numpy, torch, matplotlib
21 | ```
22 |
23 | ## 📁 Structure
24 |
25 | ```
26 | .
27 | ├── fig/
28 | ├── notebook/
29 | │ └── *
30 | ├── src/
31 | │ ├── utils/
32 | │ └── *
33 | └── README.md
34 | ```
35 |
36 | - `fig`: example figures
37 | - `notebook`: tutorials in jupyter notebook format
38 | - `src`: all source code will be here, including scripts, tools
39 | - `utils`: mini tools will be here, including dgl transferring, graph plot
40 |
41 | Here are brief guides for the supported PDE simulators:
42 | - [Guide with ARCSim](#📖-guide-with-arcsim):
43 | - [Guide with Fenics](#📖-guide-with-fenics):
44 | - [Guide with SU2](#📖-guide-with-su2):
45 |
46 | ---
47 |
48 | ## 📖 Guide with ARCSim
49 | **Step 1**. [Download](https://github.com/cbhua/tool-pdeset-generator/archive/refs/heads/main.zip) or [Clone](https://github.com/cbhua/tool-pdeset-generator.git) this repository
50 |
51 | **Step 2**. _ArcSim installation_
52 |
53 | You may find the repository with fixes [here](https://github.com/kaist-silab/arcsim) with further instructions.
54 | To install it, run the following:
55 |
56 | ```sh
57 | git clone https://github.com/kaist-silab/arcsim.git && cd arcsim/
58 | sudo chmod +x install.sh && sudo ./install.sh
59 | ```
60 |
61 | At this point, you should be ready to go.
62 |
63 | **Step 3**. _ArcSim simulation and `.obj` file saving_
64 |
65 | Let's consider the flag example. In the ArcSim folder, make a new directory called data. Then run:
66 |
67 | ```sh
68 | bin/arcsim simulate conf/flag.json data/
69 | ```
70 |
71 | (you may also run `simulateoffline` if you cannot visualize on your computer
72 |
73 | When the simulation ends (we may do that with `Esc` as well) copy the `conf/flag.json` into the folder where we saved the simulation, in our case`data/` and run:
74 |
75 | ```sh
76 | bin/arcsim generate data/
77 | ```
78 |
79 | This will generate `.obj` files that we can load into Python with `pywavefront` and the `obj_to_dgl` method we provide.
80 |
81 |
82 | 
83 |
84 |
85 | *Flag simulation with adaptive remeshing*
86 |
87 | ---
88 |
89 | ## 📖 Guide with FEniCS
90 |
91 | **Step 1**. [Download](https://github.com/cbhua/tool-pdeset-generator/archive/refs/heads/main.zip) or [Clone](https://github.com/cbhua/tool-pdeset-generator.git) this repository.
92 |
93 | **Step 2**. Based on your requirement refer to the notebooks, where there are tutorials and examples. You can find all methods provided in the below list.
94 |
95 | **Step 3**. Modify the parameters to generate your own datasets.
96 |
97 | Provided methods:
98 |
99 | - Poisson process
100 | - Customize domain & Single boundary control
101 | - Square domain & Separate boundary control
102 | - Gaussian process
103 | - Customize domain & Single boundary control (support time dynamic control)
104 | - Square domain & Separate boundary control (support time dynamic control)
105 | - Squares in square domain & Separate boundary control (support time dynamic control)
106 | - Circles in circle domain & Single boundary control (support time dynamic control)
107 |
108 | Support methods will keep updating. For more detail, you can refer to the [project manager](https://github.com/cbhua/tool-pdeset-generator/projects/1).
109 |
110 |
111 | 
112 |
113 |
114 | *Gaussian process on rectangular domain*
115 |
116 | ---
117 |
118 | ## 📖 Guide with SU2
119 |
120 | **Step 1**. [Download](https://github.com/cbhua/tool-pdeset-generator/archive/refs/heads/main.zip) or [Clone](https://github.com/cbhua/tool-pdeset-generator.git) this repository
121 |
122 | **Step 2**. Install [SU2](https://su2code.github.io/download.html)
123 |
124 | **Step 3**. _Generating data_: take a look at the notebooks and the [SU2 tutorial collection](https://su2code.github.io/tutorials/home/) to get started
125 |
126 |
127 | 
128 |
129 |
130 | *Laminar viscosity of an airfoil*
131 |
132 | ---
133 |
134 | ## 📊 Examples Gallery
135 |
136 | ### T-shirt and interactions with hard bodies
137 |
138 | 
139 |
140 |
141 | ### Paper tearing
142 |
143 | 
144 |
145 |
146 | ### Possion process, square domain, single boundary control
147 |
148 | 
149 |
150 |
151 | ### Possion process, L shape domain, single boundary control
152 |
153 | 
154 |
155 |
156 | ### Possion process, circle shape domain, single boundary control
157 |
158 | 
159 |
160 |
161 |
162 | ### Gaussian process, rectangle shape domain, multi & dynamic boundary control
163 |
164 | 
165 |
166 |
167 |
168 | ## 📜 References
169 |
170 | 1. FEniCS project: https://fenicsproject.org/
171 | 2. ARCSim project: http://graphics.berkeley.edu/resources/ARCSim/
172 | 3. DGL project: https://www.dgl.ai/
173 | 4. SU2 project: https://su2code.github.io/
174 |
--------------------------------------------------------------------------------
/fig/airfoil.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/airfoil.gif
--------------------------------------------------------------------------------
/fig/flag.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/flag.gif
--------------------------------------------------------------------------------
/fig/gaussian_circle_square.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/gaussian_circle_square.gif
--------------------------------------------------------------------------------
/fig/gaussian_square.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/gaussian_square.gif
--------------------------------------------------------------------------------
/fig/gaussian_square_dynamic.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/gaussian_square_dynamic.gif
--------------------------------------------------------------------------------
/fig/possion_circle.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/possion_circle.png
--------------------------------------------------------------------------------
/fig/possion_l.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/possion_l.png
--------------------------------------------------------------------------------
/fig/possion_square.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/possion_square.png
--------------------------------------------------------------------------------
/fig/t-shirt.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/t-shirt.gif
--------------------------------------------------------------------------------
/fig/tearing.gif:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/fig/tearing.gif
--------------------------------------------------------------------------------
/src/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/src/__init__.py
--------------------------------------------------------------------------------
/src/gaussian.py:
--------------------------------------------------------------------------------
1 | from math import degrees
2 | import numpy as np
3 | from fenics import *
4 | from dolfin import *
5 | from mshr import *
6 | from dgl.data.utils import save_graphs
7 | from src.utils.to_dgl import to_dgl
8 | from src.utils.bc import rectangle
9 |
10 |
11 | def gaussian(stop:float,
12 | steps:int,
13 | mesh,
14 | f:str='0',
15 | u0:str='0',
16 | ud: str='0',
17 | path:str='data/gaussian.bin'
18 | ):
19 | ''' Create Gaussian Equation Dataset
20 |
21 | 2D Gaussian equation, dynamic process, custom domain,
22 | singel dirichlet boundary condition.
23 |
24 | Args:
25 | stop: process stop time
26 | steps: number of time step
27 | mesh: custom mesh generated by fenics mesh tools
28 | f: right part of laplace function, in cpp argument format
29 | ud: boundary condition function, in cpp argument format
30 | u0: initialization condition function, in cpp argument format
31 | path: path for saving generated dgl graph, in .bin format
32 | '''
33 | dt = stop / steps
34 | function_space = FunctionSpace(mesh, 'P', 1)
35 | u0 = Expression(u0, degree=2)
36 | ud = Expression(ud, degree=2)
37 | f = Expression(f, degree=2)
38 |
39 | bc = DirichletBC(function_space, ud, boundary)
40 |
41 | un = interpolate(u0, function_space)
42 |
43 | u = TrialFunction(function_space)
44 | v = TestFunction(function_space)
45 | F = u * v * dx + dt * dot(grad(u), grad(v)) * dx - (un + dt * f) * v * dx
46 | a, L = lhs(F), rhs(F)
47 |
48 | u = Function(function_space)
49 | t = 0
50 | graphs = []
51 | for _ in range(steps):
52 | t += dt
53 | solve(a == L, u, bc)
54 | un.assign(u)
55 | graphs.append(to_dgl(function=u, mesh=mesh))
56 |
57 | save_graphs(path, graphs)
58 |
59 |
60 | def gaussian_square(x0:float,
61 | xn:float,
62 | y0:float,
63 | yn:float,
64 | stop:float,
65 | steps:int,
66 | f:str='0',
67 | ud_top:str='0',
68 | ud_bottom:str='0',
69 | ud_left:str='0',
70 | ud_right:str='0',
71 | u0:str='0',
72 | cell_size:float=5.,
73 | tol:float=1e-4,
74 | dy:bool=False,
75 | path:str='data/gaussian_square_static.bin'
76 | ):
77 | '''Create Gaussian Equation Dataset
78 |
79 | 2D Gaussian equation, dynamic process, rectangle domain,
80 | can custom each boundary's condition, boundary, can choose
81 | static boundary condition or dynamic boundary condition.
82 |
83 | Args:
84 | x0: left boundary for x
85 | xn: right boundary for x
86 | y0: left boundary for y
87 | yn: right boundary for y
88 | stop: process stop time
89 | steps: number of time step
90 | f: right part of laplace function, in cpp argument format
91 | ud_top: boundary condition on the top of rectangle
92 | ud_bottom: boundary condition on the bottom of rectangle
93 | ud_left: boundary condition on the left of rectangle
94 | ud_right: boundary condition on the right of rectangle
95 | u0: initialization condition function, in cpp argument format
96 | cell_siez: cell size for created mesh
97 | tol: boundary bias, e.g. (x-tol, x+tol) is a boundary on x
98 | dy: if the boundary condition is dynamic
99 | path: path for saving generated dgl graph, in .bin format
100 | '''
101 | if (dy):
102 | mesh, function_space, bc = rectangle(x0,
103 | xn,
104 | y0,
105 | yn,
106 | ud_top,
107 | ud_bottom,
108 | ud_left,
109 | ud_right,
110 | cell_size,
111 | 0,
112 | tol
113 | )
114 | else:
115 | mesh, function_space, bc = rectangle(x0,
116 | xn,
117 | y0,
118 | yn,
119 | ud_top,
120 | ud_bottom,
121 | ud_left,
122 | ud_right,
123 | cell_size,
124 | tol
125 | )
126 | dt = stop / steps
127 | u0 = Expression(u0, degree=2)
128 | f = Expression(f, degree=2)
129 | un = interpolate(u0, function_space)
130 |
131 | u = TrialFunction(function_space)
132 | v = TestFunction(function_space)
133 | F = u * v * dx + dt * dot(grad(u), grad(v)) * dx - (un + dt * f) * v * dx
134 | a, L = lhs(F), rhs(F)
135 |
136 | u = Function(function_space)
137 | t = 0
138 | graphs = []
139 | for _ in range(steps):
140 | t += dt
141 | if (dy):
142 | _, bc = rectangle(x0,
143 | xn,
144 | y0,
145 | yn,
146 | ud_top,
147 | ud_bottom,
148 | ud_left,
149 | ud_right,
150 | cell_size,
151 | t,
152 | tol
153 | )
154 | solve(a == L, u, bc)
155 | un.assign(u)
156 | graphs.append(to_dgl(function=u, mesh=mesh))
157 |
158 | save_graphs(path, graphs)
159 |
160 |
161 | def boundary(x, on_boundary):
162 | return on_boundary
--------------------------------------------------------------------------------
/src/possion.py:
--------------------------------------------------------------------------------
1 | import os
2 | from dgl.convert import graph
3 | import torch
4 | import numpy as np
5 | import dgl
6 | from fenics import *
7 | from dolfin import *
8 | from mshr import *
9 | from dgl.data.utils import save_graphs
10 | from ufl.mathfunctions import Exp
11 | from src.utils.to_dgl import to_dgl
12 | from src.utils.bc import rectangle
13 |
14 |
15 | def possion(mesh,
16 | ud,
17 | f:str='0',
18 | path:str='data/possion.bin'
19 | ):
20 | ''' Create Possion Equation Dataset
21 |
22 | 2D Possion queation, static process, custom domain,
23 | single dirichlet boundary condition.
24 |
25 | Args:
26 | mesh: custom mesh generated by fenics mesh tools
27 | f: right part of laplace function, in cpp argument format
28 | ud: boundary condition function, in cpp argument format
29 | path: path for saving generated dgl graph, in .bin format~
30 | '''
31 | function_space = FunctionSpace(mesh, 'P', 1)
32 |
33 | u = TrialFunction(function_space)
34 | v = TestFunction(function_space)
35 | f = Expression(f, degree=2)
36 | ud = Expression(ud, degree=2)
37 | bc = DirichletBC(function_space, ud, boundary)
38 |
39 | a = dot(grad(u), grad(v)) * dx
40 | L = f * v * dx
41 | u = Function(function_space)
42 | solve(a == L, u, bc)
43 |
44 | graph = to_dgl(function=u, mesh=mesh)
45 | save_graphs(path, graph)
46 |
47 |
48 | def possion_square(x0:float,
49 | xn:float,
50 | y0:float,
51 | yn:float,
52 | f:str='0',
53 | ud_top:str='0',
54 | ud_bottom:str='0',
55 | ud_left:str='0',
56 | ud_right:str='0',
57 | cell_size:float=5.,
58 | tol:float=1e-4,
59 | path:str='data/possion_square.bin'
60 | ):
61 | ''' Create Possion Equation Dataset in Square Domain
62 |
63 | 2D Possion queation, static process, rectangle domain,
64 | can custom each boundary's condition.
65 |
66 | Args:
67 | x0: left boundary for x
68 | xn: right boundary for x
69 | y0: left boundary for y
70 | yn: right boundary for y
71 | f: right part of laplace function, in cpp argument format
72 | ud_top: boundary condition on the top of rectangle
73 | ud_bottom: boundary condition on the bottom of rectangle
74 | ud_left: boundary condition on the left of rectangle
75 | ud_right: boundary condition on the right of rectangle
76 | cell_siez: cell size for created mesh
77 | tol: boundary bias, e.g. (x-tol, x+tol) is a boundary on x
78 | path: path for saving generated dgl graph, in .bin format
79 | '''
80 | mesh, function_space, bc = rectangle(x0,
81 | xn,
82 | y0,
83 | yn,
84 | ud_top,
85 | ud_bottom,
86 | ud_left,
87 | ud_right,
88 | cell_size,
89 | tol
90 | )
91 |
92 | u = TrialFunction(function_space)
93 | v = TestFunction(function_space)
94 | f = Expression(f, degree=2)
95 |
96 | a = dot(grad(u), grad(v)) * dx
97 | L = f * v * dx
98 | u = Function(function_space)
99 | solve(a == L, u, bc)
100 |
101 | graph = to_dgl(function=u, mesh=mesh)
102 | save_graphs(path, graph)
103 |
104 |
105 | def boundary(x, on_boundary):
106 | return on_boundary
107 |
--------------------------------------------------------------------------------
/src/utils/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/DiffEqML/pde-dataset-generator/45616e662ccac2449fc98e09d305936a7c15731b/src/utils/__init__.py
--------------------------------------------------------------------------------
/src/utils/arcsim.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy as np
3 | import pywavefront # for loading .obj files
4 | import logging
5 | logging.getLogger().setLevel(logging.CRITICAL) # disable long output due to error ('ms' field not compatible)
6 | default_keys=['x', 'y', 'z']
7 | default_dtype=torch.float32
8 |
9 | def arcsim_to_graph(data, **kwargs):
10 | if data is None:
11 | raise RuntimeError('Please provide .obj file to create the graph.')
12 | return obj_to_graph(data, **kwargs)
13 |
14 | def obj_to_graph(file, **kwargs):
15 | '''Transforms pywavefront object into DGL graph
16 | Compatible with ArcSim .obj files
17 | '''
18 | keys = kwargs.get('keys') if 'keys' in kwargs else default_keys
19 | dtype = kwargs.get('dtype') if 'dtype' in kwargs else default_dtype
20 | obj = pywavefront.Wavefront(file, collect_faces=True) # load ArcSim .obj file
21 | vert = np.array(obj.vertices)
22 | mesh = np.array(obj.mesh_list[0].faces)
23 | # We add the sources and destination through contiguous nodes
24 | # We should also repeat by switching indexes to have bidirectional graph
25 | # BEWARE: this will have lots of duplicates, don't think DGL can deal with them natively
26 | # extend method similar to append
27 | src = []
28 | dst = []
29 | for i, j in zip([0,1,2], [1,2,0]):
30 | src.extend(mesh[:,i].tolist()); dst.extend(mesh[:,j].tolist())
31 | src.extend(mesh[:,j].tolist()); dst.extend(mesh[:,i].tolist())
32 | nodes = vert.shape[0] - 1 # this is the number of nodes, so we have to subtract 1 for DGL graph creation
33 | edges = [src, dst]
34 | values = []
35 | for k, i in zip(keys, range(len(keys))):
36 | # Spatial coordinates
37 | if k == 'x': values.append(torch.tensor(vert[:,i], dtype=dtype))
38 | elif k == 'y': values.append(torch.tensor(vert[:,i], dtype=dtype))
39 | elif k == 'z': values.append(torch.tensor(vert[:,i], dtype=dtype))
40 | else:
41 | # Field values
42 | values.append(torch.tensor(vert[:,i], dtype=dtype)) # get data corresponding to the key
43 | return [nodes, edges, keys, values]
44 |
--------------------------------------------------------------------------------
/src/utils/bc.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from fenics import *
3 | from dolfin import *
4 | from mshr import *
5 |
6 |
7 | def rectangle(x0,
8 | xn,
9 | y0,
10 | yn,
11 | ud_top,
12 | ud_bottom,
13 | ud_left,
14 | ud_right,
15 | cell_size,
16 | tol,
17 | t=-1,
18 | ms=None,
19 | fs=None):
20 | if fs is None:
21 | domain = Rectangle(Point(x0, y0), Point(xn, yn))
22 | mesh = generate_mesh(domain, cell_size)
23 | function_space = FunctionSpace(mesh, 'P', 1)
24 | else:
25 | mesh = ms
26 | function_space = fs
27 |
28 | top = YBoundary(yn, tol)
29 | bottom = YBoundary(y0, tol)
30 | left = XBoundary(x0, tol)
31 | right = XBoundary(xn, tol)
32 |
33 | boundaries = MeshFunction('size_t', mesh, mesh.topology().dim()-1)
34 | boundaries.set_all(0)
35 | top.mark(boundaries, 1)
36 | bottom.mark(boundaries, 2)
37 | left.mark(boundaries, 3)
38 | right.mark(boundaries, 4)
39 |
40 | if t >= 0:
41 | ud_top = Expression(ud_top, degree=2, t=t)
42 | ud_bottom = Expression(ud_bottom, degree=2, t=t)
43 | ud_left = Expression(ud_left, degree=2, t=t)
44 | ud_right = Expression(ud_right, degree=2, t=t)
45 | else:
46 | ud_top = Expression(ud_top, degree=2)
47 | ud_bottom = Expression(ud_bottom, degree=2)
48 | ud_left = Expression(ud_left, degree=2)
49 | ud_right = Expression(ud_right, degree=2)
50 | bc = []
51 | bc.append(DirichletBC(function_space, ud_top, boundaries, 1))
52 | bc.append(DirichletBC(function_space, ud_bottom, boundaries, 2))
53 | bc.append(DirichletBC(function_space, ud_left, boundaries, 3))
54 | bc.append(DirichletBC(function_space, ud_right, boundaries, 4))
55 |
56 | return mesh, function_space, bc
57 |
58 |
59 | def circle(x,
60 | y,
61 | r,
62 | ud,
63 | cell_size,
64 | tol,
65 | t=-1):
66 | domain = Circle(Point(x, y), r)
67 | mesh = generate_mesh(domain, cell_size)
68 | ud = Expression(ud, degree=2)
69 | function_space = FunctionSpace(mesh, 'P', 1)
70 | if t >= 0:
71 | ud = Expression(ud, degree=2, t=t)
72 | else:
73 | bc = DirichletBC(function_space, ud, boundary)
74 | return mesh, function_space, bc
75 |
76 |
77 | def multi_rectangle(num,
78 | type,
79 | x0,
80 | xn,
81 | y0,
82 | yn,
83 | ud_top,
84 | ud_bottom,
85 | ud_left,
86 | ud_right,
87 | cell_size,
88 | tol,
89 | t):
90 | '''
91 | '''
92 | domain = Rectangle(Point(0, 0), Point(0, 0))
93 | for i in range(num):
94 | if type[i] == 1:
95 | domain += Rectangle(Point(x0[i], y0[i]), Point(xn[i], yn[i]))
96 | else:
97 | domain -= Rectangle(Point(x0[i], y0[i]), Point(xn[i], yn[i]))
98 | mesh = generate_mesh(domain, cell_size)
99 | function_space = FunctionSpace(mesh, 'P', 1)
100 | boundaries = MeshFunction('size_t', mesh, mesh.topology().dim()-1)
101 | boundaries.set_all(0)
102 | bc = []
103 | for i in range(num):
104 | top = YBoundary(yn[i], tol)
105 | bottom = YBoundary(y0[i], tol)
106 | left = XBoundary(x0[i], tol)
107 | right = XBoundary(xn[i], tol)
108 |
109 | top.mark(boundaries, 4 * i + 1)
110 | bottom.mark(boundaries, 4 * i + 2)
111 | left.mark(boundaries, 4 * i + 3)
112 | right.mark(boundaries, 4 * i + 4)
113 |
114 | if t >= 0:
115 | ud_top = Expression(ud_top[i], degree=2, t=t)
116 | ud_bottom = Expression(ud_bottom[i], degree=2, t=t)
117 | ud_left = Expression(ud_left[i], degree=2, t=t)
118 | ud_right = Expression(ud_right[i], degree=2, t=t)
119 | else:
120 | ud_top = Expression(ud_top[i], degree=2)
121 | ud_bottom = Expression(ud_bottom[i], degree=2)
122 | ud_left = Expression(ud_left[i], degree=2)
123 | ud_right = Expression(ud_right[i], degree=2)
124 |
125 | bc.append(DirichletBC(function_space, ud_top, boundaries, 4 * i + 1))
126 | bc.append(DirichletBC(function_space, ud_bottom, boundaries, 4 * i + 2))
127 | bc.append(DirichletBC(function_space, ud_left, boundaries, 4 * i + 3))
128 | bc.append(DirichletBC(function_space, ud_right, boundaries, 4 * i + 4))
129 |
130 | return mesh, function_space, bc
131 |
132 |
133 | def multi_circle(num,
134 | type,
135 | x,
136 | y,
137 | r,
138 | ud,
139 | cell_size,
140 | tol,
141 | t):
142 | domain = Circle(Point(0, 0), 0)
143 | for i in range(num):
144 | if type[i] == 1:
145 | domain += Circle(Point(x[i], y[i]), r[i])
146 | else:
147 | domain -= Circle(Point(x[i], y[i]), r[i])
148 | mesh = generate_mesh(domain, cell_size)
149 | function_space = FunctionSpace(mesh, 'P', 1)
150 | boundaries = MeshFunction('size_t', mesh, mesh.topology().dim()-1)
151 | boundaries.set_all(0)
152 | bc = []
153 | for i in range(num):
154 | round = CircleBoundary(x[i], y[i], r[i], tol)
155 | round.mark(boundaries, i)
156 | if t >= 0:
157 | ud_temp = Expression(ud[i], degree=2, t=t)
158 | else:
159 | ud_temp = Expression(ud[i], degree=2)
160 | bc.append(DirichletBC(function_space, ud_temp, boundaries, i))
161 |
162 | return mesh, function_space, bc
163 |
164 |
165 | def boundary(x, on_boundary):
166 | return on_boundary
167 |
168 |
169 | class CircleBoundary(SubDomain):
170 | def __init__(self, x, y, r, tol):
171 | SubDomain.__init__(self)
172 | self.x = x
173 | self.y = y
174 | self.r = r
175 | self.tol = tol
176 | def inside(self, x, on_boundary):
177 | flag = np.linalg.norm(x - [self.x, self.y])
178 | return near(flag, self.r, self.tol)
179 |
180 |
181 | class XBoundary(SubDomain):
182 | def __init__(self, value, tol):
183 | SubDomain.__init__(self)
184 | self.value = value
185 | self.tol = tol
186 | def inside(self, x, on_boundary):
187 | return near(x[0], self.value, self.tol)
188 |
189 |
190 | class YBoundary(SubDomain):
191 | def __init__(self, value, tol):
192 | SubDomain.__init__(self)
193 | self.value = value
194 | self.tol = tol
195 | def inside(self, x, on_boundary):
196 | return near(x[1], self.value, self.tol)
--------------------------------------------------------------------------------
/src/utils/fenics.py:
--------------------------------------------------------------------------------
1 | import torch
2 | import numpy as np
3 | from fenics import *
4 | from dolfin import *
5 | from mshr import *
6 |
7 | default_keys = ['x', 'y', 'value']
8 | # Graph Element Networks default keys
9 | gen_keys = ['x', 'y', 'coords', 'value', 'feat',
10 | 'init_feat', 'is_bdd', 'type', 'type_onehot',
11 | 'dist', 'feat', 'init_feat']
12 | default_dtype=torch.float32
13 |
14 |
15 | def fenics_to_graph(mesh, function, **kwargs):
16 | '''
17 | Generate graph values from Fenics data
18 | '''
19 | if mesh is None or function is None:
20 | raise RuntimeError('Please provide Fenics mesh and function to create the graph.')
21 | # Kwargs parsing
22 | keys = kwargs.get('keys') if 'keys' in kwargs else default_keys
23 | dtype = kwargs.get('dtype') if 'dtype' in kwargs else default_dtype
24 | use_gen = kwargs.get('use_gen') if 'use_gen' in kwargs else False
25 | # Graph creation
26 | nodes = mesh.num_vertices()
27 | src, dst = get_edges(mesh)
28 | edges = [src, dst]
29 | if use_gen is True:
30 | keys = gen_keys
31 | values = _get_gen_values(function, mesh, src, dst)
32 | else:
33 | values = []
34 | for k, i in zip(keys, range(len(keys))):
35 | # Spatial coordinates
36 | if k == 'x': values.append(torch.tensor(mesh.coordinates()[:,0], dtype=dtype))
37 | elif k == 'y': values.append(torch.tensor(mesh.coordinates()[:,1], dtype=dtype))
38 | elif k == 'z': values.append(torch.tensor(mesh.coordinates()[:,2], dtype=dtype))
39 | else:
40 | # Field values
41 | values.append(torch.tensor(get_values(function, mesh), dtype=dtype)) # get data corresponding to the key
42 | return [nodes, edges, keys, values]
43 |
44 |
45 |
46 | def get_edges(mesh):
47 | ''' Get edges for graph from mesh
48 |
49 | Get two connectivity lists of nodes, edges are directed from
50 | src to dst, required by DGL in the graph creation.
51 |
52 | Example:
53 | Edge 0->0, 0->1, 1->2 will be returned by [0, 0, 1], [0, 1, 2]
54 |
55 | Args:
56 | mesh: dolfin mesh
57 |
58 | Returns:
59 | src: [edge number] source node indexes for edges
60 | dst: [edge number] distinate node indexes for edges
61 | '''
62 | mesh.init(0,1)
63 | src = []
64 | dst = []
65 | for v in vertices(mesh):
66 | idx = v.index()
67 | neighbors = [Edge(mesh, i).entities(0) for i in v.entities(1)]
68 | neighbors = np.array(neighbors).flatten()
69 | neighbors = neighbors[np.where(neighbors != idx)[0]]
70 | for n in neighbors:
71 | src.append(int(idx))
72 | dst.append(int(n))
73 | return src, dst
74 |
75 |
76 | def get_values(u, mesh):
77 | ''' Get values for each mesh node
78 |
79 | Provid a function based on coordinate for mesh nodes, this can
80 | calculate values in each node by index order. Here the function
81 | is mostly get by FEniCS solver.
82 |
83 | Args:
84 | u: function based on
85 | nodes' coordinates
86 | mesh: dolfin mesh
87 |
88 | Returns:
89 | values: [node number] calculated result by index order
90 | '''
91 | values = []
92 | for pos in mesh.coordinates():
93 | values.append(u(Point(pos)))
94 | return values
95 |
96 |
97 | def _get_gen_values(u, mesh, src, dst):
98 | '''Get the values for the GEN model (Graph Element Networks)
99 | The order is according to the default GEN keys'''
100 |
101 | coords = mesh.coordinates()
102 | value = np.array(get_values(u, mesh)).reshape(-1, 1)
103 | feature = np.hstack([coords, value])
104 | zero_array = np.zeros((mesh.num_vertices()))
105 | one_array = np.ones((mesh.num_vertices()))
106 | onehot_array = np.hstack([zero_array.reshape(-1, 1), one_array.reshape(-1, 1)])
107 |
108 | dist = []
109 | for i in range(len(src)):
110 | dist.append(np.linalg.norm(coords[dst[i]] - coords[src[i]]))
111 | dist = np.array(dist)
112 | values = []
113 | values.append(torch.tensor(coords[:,0], dtype=torch.float32))
114 | values.append(torch.tensor(coords[:,1], dtype=torch.float32))
115 | values.append(torch.tensor(coords, dtype=torch.float32))
116 | # WARNNING: value shape here is different from original one
117 | values.append(torch.tensor(value, dtype=torch.float32))
118 | values.append(torch.tensor(feature, dtype=torch.float32))
119 | values.append(torch.tensor(feature, dtype=torch.float32))
120 | values.append(torch.tensor(zero_array))
121 | values.append(torch.tensor(one_array))
122 | values.append(torch.tensor(onehot_array))
123 | values.append(torch.tensor(dist.reshape(-1, 1), dtype=torch.float32))
124 | values.append(torch.tensor(dist.reshape(-1, 1), dtype=torch.float32))
125 | values.append(torch.tensor(dist.reshape(-1, 1), dtype=torch.float32))
126 | return values
127 |
--------------------------------------------------------------------------------
/src/utils/plot.py:
--------------------------------------------------------------------------------
1 | import os
2 | import tqdm
3 | import imageio
4 | import torch
5 | import matplotlib
6 | import matplotlib.pyplot as plt
7 | from mpl_toolkits.axes_grid1 import make_axes_locatable
8 |
9 |
10 | def plot_graph(graph,
11 | mesh_color='black',
12 | mesh_alpha=0.5,
13 | mesh_linewidth=1,
14 | contour_level=30,
15 | separate_mesh=True,
16 | figsize=(10, 5),
17 | *args,
18 | **kwargs):
19 | ''' Draw DGLGraph with node value
20 |
21 | Generate two figures, one will show the mesh, another
22 | one will show the value with node coordinates matched.
23 |
24 | Args:
25 | graph: should have node features: 'x',
26 | 'y', and 'value' and edges.
27 | mesh_color: pyplot color option
28 | mesh_alpha: pyplot alpha option
29 | mesh_linewidth: pyplot linewidht option
30 | contour_level: number of classified level in
31 | contour graph
32 | args: 2 the boundary of color bar, the first
33 | value is the lowest value, and the second value
34 | is the hightest value
35 | kwargs: arguments to be passed to matplotlib.pyplot
36 | '''
37 | plt.figure(figsize=figsize)
38 | plt.subplot(1, 2, 1)
39 | x = graph.ndata['x']
40 | y = graph.ndata['y']
41 | value = graph.ndata['value']
42 | plt.scatter(x, y, value)
43 | src, dst = graph.edges()
44 | for i in range(len(dst)):
45 | nodes_x = [x[src[i]], x[dst[i]]]
46 | nodes_y = [y[src[i]], y[dst[i]]]
47 | plt.plot(nodes_x,
48 | nodes_y,
49 | color=mesh_color,
50 | alpha=mesh_alpha,
51 | linewidth=mesh_linewidth)
52 | # Apply norm
53 | norm = None
54 | if args:
55 | norm = matplotlib.colors.Normalize(vmin=args[0], vmax=args[1])
56 | # Mesh on plot or separated
57 | cax = None
58 | if separate_mesh:
59 | ax = plt.subplot(1, 2, 2)
60 | fig = plt.tricontourf(x, y, value, levels=contour_level, norm=norm, **kwargs)
61 | divider = make_axes_locatable(ax)
62 | cax = divider.append_axes("right", size="5%", pad=0.05)
63 | # Plot with interpolation
64 | fig = plt.tricontourf(x, y, value, levels=contour_level, norm=norm, **kwargs)
65 | plt.colorbar(fig, cax=cax)
66 |
67 | def gif_generator(load_path, save_path):
68 | ''' Generate gif from series figures
69 |
70 | Args:
71 | load_path: folder contain figures in png format
72 | save_path: generated gif saved place
73 | '''
74 | with imageio.get_writer(save_path, mode='I') as writer:
75 | for root, dirs, files in os.walk(load_path):
76 | for file in tqdm.tqdm(files):
77 | image = imageio.imread(os.path.join(root, file), '.png')
78 | writer.append_data(image)
79 |
--------------------------------------------------------------------------------
/src/utils/su2.py:
--------------------------------------------------------------------------------
1 | # References: https://github.com/locuslab/cfd-gcn/blob/master/mesh_utils.py
2 | import torch
3 | import numpy as np
4 | from os import PathLike
5 | from typing import Sequence, Dict, Union, Tuple, List
6 | import vtk
7 | from vtk.numpy_interface import dataset_adapter as dsa
8 |
9 | UnionTensor = Union[torch.Tensor, np.ndarray]
10 |
11 | SU2_SHAPE_IDS = {
12 | 'line': 3,
13 | 'triangle': 5,
14 | 'quad': 9,
15 | }
16 |
17 | default_dtype = torch.float32
18 | keys = ['x', 'y', 'z']
19 |
20 | def su2_to_graph(data, mesh, **kwargs):
21 | if data is None:
22 | raise RuntimeError('Please provide .vtu or .vtk data file to create the graph.')
23 | if mesh is None:
24 | raise RuntimeError('Please provide .su2 mesh file to create the graph.')
25 | dtype = kwargs.get('dtype') if 'dtype' in kwargs else default_dtype
26 | nodes, edges, _, _ = get_mesh_graph(mesh)
27 | nodes = len(nodes)
28 | # Read the source file
29 | reader = vtk.vtkXMLUnstructuredGridReader()
30 | reader.SetFileName(data)
31 | reader.Update()
32 | output = reader.GetOutput()
33 | out = dsa.WrapDataObject(output)
34 | if 'keys' in kwargs:
35 | keys = kwargs.get('keys')
36 | else:
37 | keys = ['x', 'y', 'z'] # default spatial coordinates
38 | keys.extend(out.PointData.keys()) # read the fields
39 | values = []
40 | for k, i in zip(keys, range(len(keys))):
41 | # Spatial coordinates
42 | if k == 'x': values.append(torch.tensor(out.Points[:, 0], dtype=dtype))
43 | elif k == 'y': values.append(torch.tensor(out.Points[:, 1], dtype=dtype))
44 | elif k == 'z': values.append(torch.tensor(out.Points[:, 2], dtype=dtype))
45 | else:
46 | # Field values
47 | values.append(torch.tensor(out.PointData[k], dtype=dtype)) # get data corresponding to the key
48 | return [nodes, edges, keys, values]
49 |
50 | def get_mesh_graph(mesh_filename: Union[str, PathLike],
51 | dtype: np.dtype = np.float32
52 | ) -> Tuple[np.ndarray, np.ndarray, List[List[List[int]]], Dict[str, List[List[int]]]]:
53 |
54 | def get_rhs(s: str) -> str:
55 | return s.split('=')[-1]
56 |
57 | marker_dict = {}
58 | with open(mesh_filename) as f:
59 | for line in f:
60 | if line.startswith('NPOIN'):
61 | num_points = int(get_rhs(line))
62 | mesh_points = [[float(p) for p in f.readline().split()[:2]]
63 | for _ in range(num_points)]
64 | nodes = np.array(mesh_points, dtype=dtype)
65 |
66 | if line.startswith('NMARK'):
67 | num_markers = int(get_rhs(line))
68 | for _ in range(num_markers):
69 | line = f.readline()
70 | assert line.startswith('MARKER_TAG')
71 | marker_tag = get_rhs(line).strip()
72 | num_elems = int(get_rhs(f.readline()))
73 | marker_elems = [[int(e) for e in f.readline().split()[-2:]]
74 | for _ in range(num_elems)]
75 | # marker_dict[marker_tag] = np.array(marker_elems, dtype=np.long).transpose()
76 | marker_dict[marker_tag] = marker_elems
77 |
78 | if line.startswith('NELEM'):
79 | edges = []
80 | triangles = []
81 | quads = []
82 | num_edges = int(get_rhs(line))
83 | for _ in range(num_edges):
84 | elem = [int(p) for p in f.readline().split()]
85 | if elem[0] == SU2_SHAPE_IDS['triangle']:
86 | n = 3
87 | triangles.append(elem[1:1+n])
88 | elif elem[0] == SU2_SHAPE_IDS['quad']:
89 | n = 4
90 | quads.append(elem[1:1+n])
91 | else:
92 | raise NotImplementedError
93 | elem = elem[1:1+n]
94 | edges += [[elem[i], elem[(i+1) % n]] for i in range(n)]
95 | edges = np.array(edges, dtype=np.int32).transpose()
96 | # triangles = np.array(triangles, dtype=np.long)
97 | # quads = np.array(quads, dtype=np.long)
98 | elems = [triangles, quads]
99 |
100 | return nodes, edges, elems, marker_dict
--------------------------------------------------------------------------------
/src/utils/to_dgl.py:
--------------------------------------------------------------------------------
1 | import dgl
2 | import torch
3 | import numpy as np
4 |
5 | '''
6 | The SIMULATOR_to_graph methods return the following values:
7 | - nodes: int, number of nodes of the graph
8 | - edges: list, format [source, destination]
9 | - keys: list of strings
10 | - values: list of tensors, corresponding to the keysfrom .arcsim import arcsim_to_graph
11 | '''
12 | from .arcsim import arcsim_to_graph
13 | from .fenics import fenics_to_graph
14 | from .su2 import su2_to_graph
15 |
16 | def to_dgl(data=None,
17 | mesh=None,
18 | function=None,
19 | method=None,
20 | **kwargs):
21 | '''Obtain DGL graphs from different methods
22 | Supported file types:
23 | - arcsim (.obj)
24 | - su2 (.su2, .vtu, .vtk)
25 | If the method is not provided, we check automatically
26 | Example kwargs:
27 | - keys: list of specific field keys to look for
28 | - gen_mode: use graph element networks
29 | '''
30 | # Check for method override
31 | if method is not None:
32 | if method == 'arcsim':
33 | graph_data = arcsim_to_graph(data, **kwargs)
34 | elif method == 'fenics':
35 | graph_data = fenics_to_graph(mesh, function, **kwargs)
36 | elif method == 'su2':
37 | graph_data = su2_to_graph(data, mesh, **kwargs)
38 | else:
39 | raise ValueError("Method not supported. Currently available methods are arcsim, fenics and su2.")
40 | elif data is not None:
41 | # Automatic file detection
42 | import os
43 | filename, file_extension = os.path.splitext(data)
44 | if file_extension == '.obj':
45 | graph_data = arcsim_to_graph(data, **kwargs)
46 | elif file_extension == '.vtu' or file_extension == '.vtk':
47 | graph_data = su2_to_graph(data, mesh, **kwargs)
48 | else:
49 | raise ValueError('Detected file type is {}. Currently readable types for graph data are .obj, .bin, .vtu and .vtk').format(file_extension)
50 | elif mesh is not None and function is not None:
51 | graph_data = fenics_to_graph(mesh, function, **kwargs)
52 | else:
53 | raise ValueError('No data provided!')
54 | # Graph is in a common for all the methods
55 | nodes, edges, keys, values = graph_data
56 | return create_dgl_graph(nodes, edges, keys, values)
57 |
58 |
59 | def create_dgl_graph(nodes, edges, keys, values):
60 | '''Minimalistic approach to graph creation in DGL'''
61 | graph = dgl.DGLGraph()
62 | graph.add_nodes(nodes)
63 | graph.add_edges(edges[0], edges[1]) # edges: [sources, destinations]
64 | for k, i in zip(keys, range(len(keys))):
65 | graph.ndata[k] = values[i]
66 | return graph
67 |
68 |
--------------------------------------------------------------------------------