1+ """
2+ Visualizing computational fluid dynamics on a car
3+ ===================================================
4+ In this example we visualize a mesh drawn from the :ref:`CarCFDDataset <car_cfd_dataset_api>`.
5+ """
6+
7+ # %%
8+ # Import dependencies
9+ # --------------------
10+ # We first import our `neuralop` library and required dependencies.
11+ import numpy as np
12+ import torch
13+ import matplotlib
14+ import matplotlib .pyplot as plt
15+ from neuralop .data .datasets import load_mini_car
16+
17+ font = {'size' : 12 }
18+ matplotlib .rc ('font' , ** font )
19+
20+ torch .manual_seed (0 )
21+ np .random .seed (0 )
22+
23+ # %%
24+ # Understanding the data
25+ # ----------------------
26+ # The data in a ``MeshDataModule`` is structured as a dictionary of tensors and important scalar values encoding
27+ # a 3-d triangle mesh over the surface of a car.
28+ # Each sample includes the coordinates of all triangle vertices and the centroids of each triangle face.
29+ #
30+ # In this case, the creators used OpenFOAM to simulate the surface air pressure on car geometries in a wind tunnel.
31+ # The 3-d Navier-Stokes equations were simulated for a variety of inlet velocities over each surface using the
32+ # **OpenFOAM** computational solver to predict pressure at every vertex on the mesh.
33+ # Each sample here also has an inlet velocity scalar and a pressure field that maps 1-to-1 with the vertices on the mesh.
34+ # The :ref:`full CarCFDDataset <car_cfd_dataset_api>` is stored in triangle mesh files for downstream processing.
35+ # For the sake of simplicity, we've packaged a few examples of the data after processing in tensor form to visualize here:
36+
37+ data_list = load_mini_car ()
38+ sample = data_list [0 ]
39+ print (f'{ sample .keys ()= } ' )
40+
41+ # %%
42+ # Visualizing the car
43+ # -------------------
44+ # Let's take a look at the vertices and pressure values.
45+
46+ fig = plt .figure ()
47+ ax = fig .add_subplot (projection = '3d' )
48+ # By default the data is normalized into the unit cube. To get a
49+ # better look at it, we scale the z-direction up.
50+ scatter = ax .scatter (sample ['vertices' ][:,0 ],sample ['vertices' ][:,1 ],
51+ sample ['vertices' ][:,2 ]* 2 , s = 2 , c = sample ['press' ])
52+ ax .set_xlim (0 ,2 )
53+ ax .set_ylim (0 ,2 )
54+ ax .set_xlabel ("x" )
55+ ax .set_ylabel ("y" )
56+ ax .set_zlabel ("z" )
57+ ax .view_init (elev = 20 , azim = 150 , roll = 0 , vertical_axis = 'y' )
58+ ax .set_title ("Pressure over car mesh vertices" )
59+ fig .colorbar (scatter , pad = 0.2 , label = "normalized pressure" , ax = ax )
60+ plt .draw ()
61+ # %%
62+ # Query points
63+ # -------------
64+ # Each sample in the ``CarCFDDataset`` also includes a set of latent query points on which we learn a function
65+ # to enable learning with an FNO in the middle of our geometry-informed models. Let's visualize the queries
66+ # on top of the car from before:
67+ fig = plt .figure (figsize = (8 ,10 ))
68+ ax = fig .add_subplot (projection = '3d' )
69+ scatter = ax .scatter (sample ['vertices' ][:,0 ],sample ['vertices' ][:,1 ],
70+ sample ['vertices' ][:,2 ]* 2 , s = 4 , label = 'Car surface' )
71+ queries = sample ['query_points' ].view (- 1 , 3 ) # unroll our cube tensor into a point cloud
72+ ax .scatter (queries [:,0 ],queries [:,1 ],queries [:,2 ]* 2 ,s = 1 , alpha = 0.5 , label = 'Latent queries' )
73+
74+ ax .set_xlim (0 ,2 )
75+ ax .set_ylim (0 ,2 )
76+ ax .set_xlabel ("x" )
77+ ax .set_ylabel ("y" )
78+ ax .set_zlabel ("z" )
79+ ax .legend ()
80+ ax .view_init (elev = 20 , azim = 150 , roll = 0 , vertical_axis = 'y' )
81+ ax .set_title ("Query points and vertices" )
82+ # %%
83+ # Neighbor search between 3D point clouds
84+ # In :doc:`../layers/plot_neighbor_search` we demonstrate our neighbor search
85+ # on a simple 2-d point cloud. Let's try that again with our points here:
86+
87+ from neuralop .layers .neighbor_search import native_neighbor_search
88+ verts = sample ['vertices' ]
89+ #query_point = queries[1000]
90+ query_point = queries [3300 ] # 1550 and 0.4 is really good
91+ #nbr_data = native_neighbor_search(data=verts, queries=query_point.unsqueeze(0), radius=0.15)
92+ nbr_data = native_neighbor_search (data = verts , queries = query_point .unsqueeze (0 ), radius = 0.5 )
93+
94+ # %% Visualizing neighborhoods
95+ # Let's plot the new neighbors we just found on top of the car surface point cloud.
96+ fig = plt .figure (figsize = (8 ,10 ))
97+ ax1 = fig .add_subplot (2 ,1 ,1 , projection = '3d' )
98+ ax2 = fig .add_subplot (2 ,1 ,2 , projection = '3d' )
99+ neighbors = verts [nbr_data ['neighbors_index' ]]
100+
101+ # Plotting just one query point vs. the car
102+ ax1 .scatter (verts [:, 0 ], verts [:, 1 ], verts [:, 2 ]* 2 , s = 1 , alpha = 0.1 )
103+ ax1 .scatter (query_point [0 ], query_point [1 ], query_point [2 ]* 2 , s = 10 , c = 'red' , label = 'Single query' )
104+ ax1 .view_init (elev = 20 , azim = - 20 , roll = 0 , vertical_axis = 'y' )
105+ ax1 .legend ()
106+ ax1 .set_xlim (0 ,2 )
107+ ax1 .set_ylim (0 ,2 )
108+ ax1 .set_xlabel ("x" )
109+ ax1 .set_ylabel ("y" )
110+ ax1 .set_zlabel ("z" )
111+ ax1 .view_init (elev = 20 , azim = - 20 , roll = 0 , vertical_axis = 'y' )
112+ ax1 .grid (False )
113+ ax1 .set_title ("One query point" )
114+
115+ # Plotting all query points and neighbors
116+ ax2 .scatter (verts [:, 0 ], verts [:, 1 ], verts [:, 2 ]* 2 , s = 0.5 , alpha = 0.4 , label = "Car surface" )
117+ ax2 .scatter (queries [:, 0 ], queries [:, 1 ], queries [:, 2 ]* 2 , s = 0.5 , alpha = 0.2 , label = "Latent queries" )
118+ ax2 .scatter (neighbors [:, 0 ], neighbors [:, 1 ], neighbors [:, 2 ]* 2 , s = 10 , label = "Neighbors on\n car surface" ,)
119+ ax2 .legend ()
120+ ax2 .set_xlim (0 ,2 )
121+ ax2 .set_ylim (0 ,2 )
122+ ax2 .set_xlabel ("x" )
123+ ax2 .set_ylabel ("y" )
124+ ax2 .set_zlabel ("z" )
125+ ax2 .view_init (elev = 20 , azim = - 20 , roll = 0 , vertical_axis = 'y' )
126+ ax2 .set_title ("Neighbor points from car for one query point" )
127+ ax2 .grid (False )
128+
129+ for ax in ax1 ,ax2 :
130+ ax .set_xticks ([])
131+ ax .set_yticks ([])
132+ ax .set_zticks ([])
133+ plt .draw ()
134+
135+
136+ # %%
137+ # **Connecting neighbors to query**
138+ #
139+ # First, let's make a simple utiltiy to add arrows to our 3D plot:
140+
141+ import numpy as np
142+ from matplotlib import pyplot as plt
143+ from matplotlib .patches import FancyArrowPatch
144+ from mpl_toolkits .mplot3d import proj3d
145+
146+ class Arrow3D (FancyArrowPatch ):
147+ def __init__ (self , xs , ys , zs , * args , ** kwargs ):
148+ super ().__init__ ((0 ,0 ), (0 ,0 ), * args , ** kwargs )
149+ self ._verts3d = xs , ys , zs
150+
151+ def do_3d_projection (self , renderer = None ):
152+ xs3d , ys3d , zs3d = self ._verts3d
153+ xs , ys , zs = proj3d .proj_transform (xs3d , ys3d , zs3d , self .axes .M )
154+ self .set_positions ((xs [0 ],ys [0 ]),(xs [1 ],ys [1 ]))
155+
156+ return np .min (zs )
157+
158+ # Creating plots
159+ fig = plt .figure (figsize = (8 ,10 ))
160+ ax1 = fig .add_subplot (projection = '3d' )
161+ neighbors = verts [nbr_data ['neighbors_index' ]]
162+
163+ # Plotting just one query point vs. the car
164+ ax1 .scatter (verts [:, 0 ], verts [:, 1 ], verts [:, 2 ]* 2 , s = 1 , alpha = 0.1 )
165+ ax1 .scatter (query_point [0 ], query_point [1 ], query_point [2 ]* 2 , s = 10 , c = 'red' , label = 'Single query' )
166+ ax1 .scatter (neighbors [:, 0 ], neighbors [:, 1 ], neighbors [:, 2 ]* 2 , s = 10 , label = "Neighbors on\n car surface" ,)
167+
168+ ax1 .view_init (elev = 20 , azim = - 20 , roll = 0 , vertical_axis = 'y' )
169+ ax1 .legend ()
170+ ax1 .set_xlim (0 ,2 )
171+ ax1 .set_ylim (0 ,2 )
172+ ax1 .set_xlabel ("x" )
173+ ax1 .set_ylabel ("y" )
174+ ax1 .set_zlabel ("z" )
175+ ax1 .view_init (elev = 20 , azim = - 20 , roll = 0 , vertical_axis = 'y' )
176+ ax1 .grid (False )
177+ ax1 .set_title ("One query point" )
178+
179+
180+ for ax in [ax1 ]:
181+ ax .set_xticks ([])
182+ ax .set_yticks ([])
183+ ax .set_zticks ([])
184+
185+ # Add arrows between neighbors and query
186+ arrow_prop_dict = dict (mutation_scale = 1 , arrowstyle = '-|>' , color = 'red' , shrinkA = 1 , shrinkB = 1 , alpha = 0.1 )
187+ for nbr in neighbors :
188+ a = Arrow3D ([query_point [0 ], nbr [0 ]],
189+ [query_point [1 ], nbr [1 ]],
190+ [query_point [2 ]* 2 , nbr [2 ]* 2 ], ** arrow_prop_dict )
191+ ax1 .add_artist (a )
192+
193+ fig .tight_layout ()
194+ plt .draw ()
0 commit comments