├── .gitattributes ├── .gitignore ├── CITATION.cff ├── LICENSE ├── README.md ├── app └── interactive_app.py ├── data └── test.png ├── media └── demo2.gif ├── neuralff ├── __init__.py ├── field.py ├── loss │ └── __init__.py ├── model │ ├── BasicNetwork.py │ └── __init__.py └── ops │ ├── __init__.py │ ├── fluid_ops.py │ ├── grid_ops.py │ ├── physics_ops.py │ └── vector_ops.py └── setup.py /.gitattributes: -------------------------------------------------------------------------------- 1 | *.png filter=lfs diff=lfs merge=lfs -text 2 | *.gif filter=lfs diff=lfs merge=lfs -text 3 | *.jpeg filter=lfs diff=lfs merge=lfs -text 4 | *.jpg filter=lfs diff=lfs merge=lfs -text 5 | *.PNG filter=lfs diff=lfs merge=lfs -text 6 | *.GIF filter=lfs diff=lfs merge=lfs -text 7 | *.JPEG filter=lfs diff=lfs merge=lfs -text 8 | *.JPG filter=lfs diff=lfs merge=lfs -text 9 | *.pdf filter=lfs diff=lfs merge=lfs -text 10 | *.PDF filter=lfs diff=lfs merge=lfs -text 11 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.sw* 2 | __pycache__ 3 | *.egg-info 4 | 5 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: "If you use this software, please cite it as below." 3 | authors: 4 | - family-names: "Takikawa" 5 | given-names: "Towaki" 6 | orcid: "https://orcid.org/0000-0003-2019-1564" 7 | - family-names: "Mozaffari" 8 | given-names: "Mohammad" 9 | orcid: "https://orcid.org/0000-0000-0000-0000" 10 | title: "Neural Fluid Fields" 11 | version: 1.0.0 12 | date-released: 2021-12-07 13 | url: "https://github.com/tovacinni/neural-fluid-fields" 14 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2021 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | this software and associated documentation files (the "Software"), to deal in 7 | the Software without restriction, including without limitation the rights to 8 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | the Software, and to permit persons to whom the Software is furnished to do so, 10 | subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Neural Fluid Fields 2 | 3 | drawing 4 | 5 | This is a small library for doing fluid simulation with neural fields. 6 | Check out our review paper, [Neural Fields in Visual Computing and Beyond](https://neuralfields.cs.brown.edu/) 7 | if you want to learn more about neural fields! 8 | 9 | ## Code Organization 10 | 11 | `neuralff` contains the bulk of the library. The library can be installed as a usual Python module, 12 | by doing `python setup.py develop` or any equivalent. 13 | 14 | The library contains of several components: 15 | 16 | The `neuralff.ops` module contains the core utility functions. In particular, 17 | `neuralff/ops/fluid_physics_ops.py` contains PDE loss functions, `neuralff/ops/vector_ops.py` contains 18 | differential operators, and `neuralff/ops/fluid_ops.py` contains functions for performing advection. 19 | 20 | These functions generally take as input a `neuralff.Field` class, which can be a grid-based vector field, 21 | or a neural field so long as it can sample vectors on continuous coordinates with some mechanism. 22 | 23 | ## Running the Demo 24 | 25 | The demo is located in `app`. These are standalone demos which use the `neuralff` library to do things like 26 | real-time fluid simulation using neural fields. 27 | 28 | The demo runs on `glumpy` and `pycuda`, which can be annoying to install. To install: 29 | 30 | ``` 31 | git clone https://github.com/inducer/pycuda 32 | git submodule update --recursive --init 33 | python configure.py --cuda-root=$CUDA_HOME --cuda-enable-gl 34 | python setup.py develop 35 | pip install pyopengl 36 | pip install glumpy 37 | ``` 38 | 39 | To run the demo, simply run `python3 app/interactive_app.py`. 40 | 41 | To change the inital image of the animation, set an arbitrary image named `test.png` in `./data`. 42 | 43 | 44 | -------------------------------------------------------------------------------- /app/interactive_app.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import os 23 | import sys 24 | from contextlib import contextmanager 25 | import importlib 26 | 27 | import numpy as np 28 | import torch 29 | import torch.nn as nn 30 | import torch.nn.functional as F 31 | import torch.optim as optim 32 | 33 | import pycuda.driver 34 | from pycuda.gl import graphics_map_flags 35 | from glumpy import app, gloo, gl 36 | 37 | from torch.profiler import profile, record_function, ProfilerActivity 38 | 39 | import neuralff 40 | from neuralff.model import BasicNetwork 41 | import neuralff.ops as nff_ops 42 | 43 | from scipy import sparse 44 | import scipy.sparse.linalg as linalg 45 | 46 | import cv2 47 | import skimage 48 | import imageio 49 | 50 | import tqdm 51 | 52 | import argparse 53 | 54 | def load_rgb(path): 55 | img = imageio.imread(path) 56 | img = skimage.img_as_float32(img) 57 | img = img[:,:,:3] 58 | #img -= 0.5 59 | #img *= 2. 60 | #img = img.transpose(2, 0, 1) 61 | img = img.transpose(1, 0, 2) 62 | return img 63 | 64 | @contextmanager 65 | def cuda_activate(img): 66 | """Context manager simplifying use of pycuda.gl.RegisteredImage""" 67 | mapping = img.map() 68 | yield mapping.array(0,0) 69 | mapping.unmap() 70 | 71 | def create_shared_texture(w, h, c=4, 72 | map_flags=graphics_map_flags.WRITE_DISCARD, 73 | dtype=np.uint8): 74 | """Create and return a Texture2D with gloo and pycuda views.""" 75 | tex = np.zeros((h,w,c), dtype).view(gloo.Texture2D) 76 | tex.activate() # force gloo to create on GPU 77 | tex.deactivate() 78 | cuda_buffer = pycuda.gl.RegisteredImage(int(tex.handle), tex.target, map_flags) 79 | return tex, cuda_buffer 80 | 81 | backend = "glumpy.app.window.backends.backend_glfw" 82 | importlib.import_module(backend) 83 | 84 | def load_rgb(path): 85 | img = imageio.imread(path) 86 | img = skimage.img_as_float32(img) 87 | img = img[:,:,:3] 88 | return img 89 | 90 | def resize_rgb(img, height, width, interpolation=cv2.INTER_LINEAR): 91 | img = cv2.resize(img, dsize=(height, width), interpolation=interpolation) 92 | return img 93 | 94 | class InteractiveApp(sys.modules[backend].Window): 95 | 96 | #def __init__(self, render_res=[720, 1024]): 97 | #def __init__(self, render_res=[100, 200]): 98 | def __init__(self, args): 99 | 100 | self.args = args 101 | self.rgb = torch.from_numpy(load_rgb(self.args.image_path)).cuda() 102 | render_res = self.rgb.shape[:2] 103 | self.render_res = render_res 104 | 105 | print("Controls:") 106 | print("h,l: switch optimization modes") 107 | print("j,k: switch display buffer") 108 | print("q : quit simulation") 109 | print("n : begin simulation") 110 | 111 | super().__init__(width=render_res[1], height=render_res[0], 112 | fullscreen=False, config=app.configuration.get_default()) 113 | 114 | import pycuda.gl.autoinit 115 | import pycuda.gl 116 | assert torch.cuda.is_available() 117 | print('using GPU {}'.format(torch.cuda.current_device())) 118 | self.buffer = torch.zeros(*render_res, 4, device='cuda') 119 | 120 | self.camera_origin = np.array([2.5, 2.5, 2.5]) 121 | self.world_transform = np.eye(3) 122 | 123 | self.tex, self.cuda_buffer = create_shared_texture(self.width, self.height, 4) 124 | 125 | vertex = """ 126 | uniform float scale; 127 | attribute vec2 position; 128 | attribute vec2 texcoord; 129 | varying vec2 v_texcoord; 130 | void main() 131 | { 132 | v_texcoord = texcoord; 133 | gl_Position = vec4(scale*position, 0.0, 1.0); 134 | } """ 135 | 136 | fragment = """ 137 | uniform sampler2D tex; 138 | varying vec2 v_texcoord; 139 | void main() 140 | { 141 | gl_FragColor = texture2D(tex, v_texcoord); 142 | } """ 143 | 144 | self.screen = gloo.Program(vertex, fragment, count=4) 145 | self.screen['position'] = [(-1,-1), (-1,+1), (+1,-1), (+1,+1)] 146 | self.screen['texcoord'] = [(0,0), (0,1), (1,0), (1,1)] 147 | self.screen['scale'] = 1.0 148 | self.screen['tex'] = self.tex 149 | 150 | #self.mode = "stable_fluids" 151 | self.mode = "neuralff" 152 | 153 | self.display_modes = ["rgb", "pressure", "velocity", "rho", "divergence", "euler"] 154 | self.display_mode_idx = 0 155 | self.display_mode = self.display_modes[self.display_mode_idx] 156 | 157 | self.optim_modes = ["euler", "divergence-free", "split"] 158 | self.optim_mode_idx = 0 159 | self.optim_mode = self.optim_modes[self.optim_mode_idx] 160 | 161 | self.max_euler_error = 0.0 162 | self.max_divergence_error = 0.0 163 | self.curr_error = 0.0 164 | self.optim_switch = False 165 | self.begin_switch = False 166 | 167 | def on_draw(self, dt): 168 | title = f"FPS: {self.fps:.3f}" 169 | title += f" Buffer: {self.display_mode}" 170 | 171 | if self.display_mode == "divergence" or self.display_mode == "euler": 172 | title += f" Error: {self.curr_error:.3e}" 173 | 174 | title += f" Optimizing: {self.optim_mode}" 175 | 176 | self.set_title(title.encode("ascii")) 177 | tex = self.screen['tex'] 178 | h,w = tex.shape[:2] 179 | 180 | # render with pytorch 181 | state = torch.zeros(*self.render_res, 4, device='cuda') 182 | 183 | coords = nff_ops.normalized_grid_coords(*self.render_res) 184 | 185 | out = self.render(coords) 186 | 187 | write_dim = out.shape[-1] 188 | 189 | state[...,:write_dim] = out 190 | state[...,3] = 1 191 | state = torch.flip(state, [0]) 192 | 193 | img = (255*state).byte().contiguous() 194 | 195 | # copy from torch into buffer 196 | assert tex.nbytes == img.numel()*img.element_size() 197 | with cuda_activate(self.cuda_buffer) as ary: 198 | cpy = pycuda.driver.Memcpy2D() 199 | cpy.set_src_device(img.data_ptr()) 200 | cpy.set_dst_array(ary) 201 | cpy.width_in_bytes = cpy.src_pitch = cpy.dst_pitch = tex.nbytes//h 202 | cpy.height = h 203 | cpy(aligned=False) 204 | torch.cuda.synchronize() 205 | # draw to screen 206 | self.clear() 207 | self.screen.draw(gl.GL_TRIANGLE_STRIP) 208 | 209 | def on_close(self): 210 | pycuda.gl.autoinit.context.pop() 211 | 212 | #################################### 213 | # Application specific code 214 | #################################### 215 | 216 | def on_key_press(self, symbol, modifiers): 217 | if symbol == 75: # k 218 | self.display_mode_idx = (self.display_mode_idx + 1) % len(self.display_modes) 219 | self.display_mode = self.display_modes[self.display_mode_idx] 220 | elif symbol == 74: # j 221 | self.display_mode_idx = (self.display_mode_idx - 1) % len(self.display_modes) 222 | self.display_mode = self.display_modes[self.display_mode_idx] 223 | elif symbol == 81: # q 224 | self.close() 225 | elif symbol == 78: # n 226 | self.optim_switch = not self.optim_switch 227 | elif symbol == 76: # l 228 | self.optim_mode_idx = (self.optim_mode_idx + 1) % len(self.optim_modes) 229 | self.optim_mode = self.optim_modes[self.optim_mode_idx] 230 | elif symbol == 72: # h 231 | self.optim_mode_idx = (self.optim_mode_idx - 1) % len(self.optim_modes) 232 | self.optim_mode = self.optim_modes[self.optim_mode_idx] 233 | elif symbol == 66: # b 234 | self.begin_switch = not self.begin_switch 235 | 236 | def init_state(self): 237 | 238 | #self.gravity = 1e-4 239 | #self.timestep = 1e-1 240 | #self.timestep = 5e-2 241 | #self.timestep = 1e-5 242 | 243 | #self.timestep = 5e-3 244 | #self.timestep = 5e-2 245 | self.timestep = 1e-1 246 | 247 | self.image_coords = nff_ops.normalized_grid_coords(self.height, self.width, aspect=False, device="cuda") 248 | self.image_coords[...,1] *= -1 249 | 250 | if self.mode == "stable_fluids": 251 | self.grid_width = self.width // 8 252 | self.grid_height = self.height // 8 253 | 254 | elif self.mode == "neuralff": 255 | 256 | velocity_field_config = { 257 | "input_dim" : 2, 258 | "output_dim" : 2, 259 | "hidden_activation" : torch.sin, 260 | "output_activation" : None, 261 | "bias" : True, 262 | "num_layers" : 4, 263 | "hidden_dim" : 128, 264 | } 265 | 266 | self.velocity_field = neuralff.NeuralField(**velocity_field_config).cuda() 267 | 268 | pressure_field_config = { 269 | "input_dim" : 2, 270 | "output_dim" : 1, 271 | "hidden_activation" : torch.sin, 272 | "output_activation" : None, 273 | "bias" : True, 274 | "num_layers" : 4, 275 | "hidden_dim" : 128, 276 | } 277 | 278 | self.pressure_field = neuralff.NeuralField(**pressure_field_config).cuda() 279 | 280 | self.rho_field = neuralff.ImageDensityField(self.height, self.width) 281 | self.rho_field.update(self.rgb) 282 | 283 | self.pc_lr = self.args.pc_lr 284 | self.precondition_optimizer = optim.Adam([ 285 | {"params": self.velocity_field.parameters(), "lr":self.pc_lr}, 286 | {"params": self.pressure_field.parameters(), "lr":self.pc_lr}, 287 | ]) 288 | 289 | self.lr = self.args.lr 290 | self.optimizer = optim.Adam([ 291 | #self.optimizer = optim.SGD([ 292 | {"params": self.velocity_field.parameters(), "lr":self.lr}, 293 | {"params": self.pressure_field.parameters(), "lr":self.lr}, 294 | ]) 295 | 296 | if self.args.precondition: 297 | self.precondition() 298 | 299 | def precondition(self): 300 | 301 | num_batch = self.args.pc_num_batch 302 | batch_size = self.args.pc_batch_size 303 | epochs = self.args.pc_epochs 304 | pts = torch.rand([batch_size*num_batch, 2], device='cuda') * 2.0 - 1.0 305 | 306 | initial_velocity = self.velocity_field.sample(pts).detach() 307 | print("Preconditioning body forces...") 308 | for i in tqdm.tqdm(range(epochs)): 309 | for j in range(num_batch): 310 | self.velocity_field.zero_grad() 311 | self.pressure_field.zero_grad() 312 | 313 | loss = nff_ops.body_forces_loss( 314 | pts[j*batch_size:(j+1)*batch_size], 315 | self.velocity_field, self.timestep, 316 | initial_velocity=initial_velocity[j*batch_size:(j+1)*batch_size]) 317 | loss = loss.mean() 318 | loss.backward() 319 | self.precondition_optimizer.step() 320 | 321 | print("Preconditioning divergence...") 322 | for i in tqdm.tqdm(range(epochs)): 323 | for j in range(num_batch): 324 | self.velocity_field.zero_grad() 325 | self.pressure_field.zero_grad() 326 | 327 | loss = nff_ops.divergence_free_loss( 328 | pts[j*batch_size:(j+1)*batch_size], 329 | self.velocity_field) 330 | loss = loss.mean() 331 | loss.backward() 332 | self.precondition_optimizer.step() 333 | 334 | initial_velocity = self.velocity_field.sample(pts).detach() 335 | 336 | print("Preconditioning Euler...") 337 | for i in tqdm.tqdm(range(epochs)): 338 | for j in range(num_batch): 339 | self.velocity_field.zero_grad() 340 | self.pressure_field.zero_grad() 341 | 342 | loss = nff_ops.euler_loss( 343 | pts[j*batch_size:(j+1)*batch_size], 344 | self.velocity_field, 345 | self.pressure_field, self.rho_field, self.timestep, 346 | initial_velocity=initial_velocity[j*batch_size:(j+1)*batch_size]) 347 | loss = loss.mean() 348 | loss.backward() 349 | self.precondition_optimizer.step() 350 | 351 | def render(self, coords): 352 | 353 | self.optimizer = optim.Adam([ 354 | #self.optimizer = optim.SGD([ 355 | {"params": self.velocity_field.parameters(), "lr":self.lr}, 356 | {"params": self.pressure_field.parameters(), "lr":self.lr}, 357 | ]) 358 | 359 | if self.mode == "stable_fluids": 360 | # Add external forces 361 | self.velocity_field.vector_field[..., 1] += 9.8 * self.timestep * self.gravity 362 | 363 | # Remove divergence 364 | #self.velocities = remove_divergence(self.velocities, self.x_mapper, self.y_mapper) 365 | 366 | elif self.optim_switch: 367 | for i in range(6): 368 | self.pressure_field.zero_grad() 369 | self.velocity_field.zero_grad() 370 | pts = torch.rand([self.args.batch_size, 2], device=coords.device) * 2.0 - 1.0 371 | if self.optim_mode == "divergence-free": 372 | loss = nff_ops.divergence_free_loss(pts, self.velocity_field) 373 | elif self.optim_mode == "split": 374 | loss = nff_ops.body_forces_loss(pts, self.velocity_field, self.timestep) +\ 375 | nff_ops.incompressibility_loss(pts, self.velocity_field, self.pressure_field, 376 | self.rho_field, self.timestep) 377 | elif self.optim_mode == "euler": 378 | loss = nff_ops.euler_loss(pts, self.velocity_field, 379 | self.pressure_field, self.rho_field, self.timestep) 380 | loss.mean().backward() 381 | self.optimizer.step() 382 | 383 | 384 | with torch.no_grad(): 385 | if self.display_mode == "rgb": 386 | if self.begin_switch: 387 | self.rho_field.update(nff_ops.semi_lagrangian_advection( 388 | self.image_coords, self.rho_field.vector_field, self.velocity_field, self.timestep)) 389 | return self.rho_field.vector_field 390 | elif self.display_mode == "pressure": 391 | return (1.0 + self.pressure_field.sample(self.image_coords)) / 2.0 392 | elif self.display_mode == "velocity": 393 | return (1.0 + F.normalize(self.velocity_field.sample(self.image_coords), dim=-1)) / 2.0 394 | elif self.display_mode == "rho": 395 | rfsample = self.rho_field.sample(self.image_coords) 396 | return rfsample / rfsample.max() 397 | elif self.display_mode == "divergence": 398 | div = nff_ops.divergence(self.image_coords, self.velocity_field, method='finitediff')**2 399 | err = div.max() 400 | self.curr_error = err 401 | self.max_divergence_error = max(err, self.max_divergence_error) 402 | return div / self.max_divergence_error 403 | #return div / err 404 | elif self.display_mode == "euler": 405 | loss = nff_ops.euler_loss(self.image_coords, self.velocity_field, 406 | self.pressure_field, self.rho_field, self.timestep) 407 | err = loss.max() 408 | self.curr_error = err 409 | self.max_euler_error = max(err, self.max_euler_error) 410 | #return loss / self.max_euler_error 411 | return loss / err 412 | else: 413 | return torch.zeros_like(coords) 414 | 415 | def parse_options(): 416 | parser = argparse.ArgumentParser(description='Fluid simulation with neural networks.') 417 | 418 | # Global arguments 419 | global_group = parser.add_argument_group('global') 420 | global_group.add_argument('--lr', type=float, default=1e-6, 421 | help='Learning rate for the simulation.') 422 | global_group.add_argument('--batch_size', type=int, default=4096, 423 | help='Batch size for the simulation.') 424 | global_group.add_argument('--pc_lr', type=float, default=1e-6, 425 | help='Learning rate for the preconditioner.') 426 | global_group.add_argument('--pc_batch_size', type=int, default=4096, 427 | help='Batch size for the preconditioner.') 428 | global_group.add_argument('--pc_num_batch', type=int, default=10, 429 | help='Number of batches to use for the preconditioner.') 430 | global_group.add_argument('--pc_epochs', type=int, default=100, 431 | help='Number of epochs to train the preconditioner for.') 432 | global_group.add_argument('--precondition', action='store_true', 433 | help='Use the preconditioner.') 434 | global_group.add_argument('--image_path', type=str, default="./data/test.png", 435 | help='Path to the image to use for the simulation.') 436 | 437 | return parser.parse_args() 438 | 439 | if __name__=='__main__': 440 | args = parse_options() 441 | app.use('glfw') 442 | window = InteractiveApp(args) 443 | window.init_state() 444 | app.run() 445 | 446 | 447 | -------------------------------------------------------------------------------- /data/test.png: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:29cf3fcfb178c9c0b948fc282020180f6c3d969b0041cc8aa0f14a1aab7ab125 3 | size 4677 4 | -------------------------------------------------------------------------------- /media/demo2.gif: -------------------------------------------------------------------------------- 1 | version https://git-lfs.github.com/spec/v1 2 | oid sha256:1d653506df931a23eed4b2bc2415c9c865b2bc0fd211ddb161edb50d9c0484e3 3 | size 7316599 4 | -------------------------------------------------------------------------------- /neuralff/__init__.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import neuralff.ops 23 | import neuralff.model 24 | import neuralff.loss 25 | from .field import * 26 | 27 | -------------------------------------------------------------------------------- /neuralff/field.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import torch 23 | import torch.nn as nn 24 | import torch.nn.functional as F 25 | 26 | from neuralff.ops import sample_from_grid 27 | from neuralff.model import BasicNetwork 28 | 29 | class BaseField(nn.Module): 30 | def sample(self, coords): 31 | raise NotImplementedError 32 | def forward(self, coords): 33 | return self.sample(coords) 34 | 35 | 36 | class RegularVectorField(BaseField): 37 | def __init__(self, height, width, fdim=2): 38 | super().__init__() 39 | self.vector_field = nn.Parameter(torch.randn([height, width, fdim])) 40 | 41 | def sample(self, coords): 42 | shape = coords.shape 43 | samples = sample_from_grid(coords.reshape(-1, shape[-1]), self.vector_field) 44 | sample_dim = samples.shape[-1] 45 | return samples.reshape(*shape[:-1], sample_dim) 46 | 47 | def update(self, vector_field): 48 | self.vector_field = nn.Parameter(vector_field) 49 | 50 | class ImageDensityField(RegularVectorField): 51 | def sample(self, coords): 52 | alpha = 3.0 - super().sample(coords).sum(-1, keepdim=True) 53 | return (alpha + 1e-1) * 255 54 | 55 | class NeuralField(BaseField): 56 | def __init__(self, **kwargs): 57 | super().__init__() 58 | self.vector_field = BasicNetwork(**kwargs) 59 | 60 | def sample(self, coords): 61 | vector = self.vector_field(coords*100) 62 | #vector[torch.abs(coords) >= 1.0] = 0 63 | return vector 64 | 65 | class RegularNeuralField(BaseField): 66 | def __init__(self, height, width, fdim, **kwargs): 67 | super().__init__() 68 | self.vector_field = BasicNetwork(**kwargs) 69 | self.feature_field = nn.Parameter(torch.randn([height, width, fdim])) 70 | 71 | def sample(self, coords): 72 | shape = coords.shape 73 | features = sample_from_grid(coords.reshape(-1, shape[-1]), self.feature_field) 74 | feature_dim = features.shape[-1] 75 | features = features.reshape(*shape[:-1], feature_dim) 76 | return self.vector_field(features) 77 | 78 | 79 | -------------------------------------------------------------------------------- /neuralff/loss/__init__.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | 23 | -------------------------------------------------------------------------------- /neuralff/model/BasicNetwork.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import torch 23 | import torch.nn as nn 24 | import torch.nn.functional as F 25 | 26 | class BasicNetwork(nn.Module): 27 | def __init__(self, 28 | input_dim = 2, 29 | output_dim = 2, 30 | hidden_activation = torch.sin, 31 | output_activation = torch.tanh, 32 | bias = True, 33 | num_layers = 5, 34 | hidden_dim = 32): 35 | 36 | super().__init__() 37 | self.hidden_activation = hidden_activation 38 | self.output_activation = output_activation 39 | layers = [] 40 | for i in range(num_layers): 41 | if i == 0: 42 | layers.append(nn.Linear(input_dim, hidden_dim, bias=bias)) 43 | else: 44 | layers.append(nn.Linear(hidden_dim, hidden_dim, bias=bias)) 45 | self.layers = nn.ModuleList(layers) 46 | self.lout = nn.Linear(hidden_dim, output_dim, bias=bias) 47 | 48 | def forward(self, x): 49 | h = x 50 | for i, l in enumerate(self.layers): 51 | if self.hidden_activation is not None: 52 | h = self.hidden_activation(l(h)) 53 | else: 54 | h = l(h) 55 | if self.output_activation is not None: 56 | out = self.output_activation(self.lout(h)) 57 | else: 58 | out = self.lout(h) 59 | return out 60 | 61 | -------------------------------------------------------------------------------- /neuralff/model/__init__.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | from .BasicNetwork import * 23 | 24 | -------------------------------------------------------------------------------- /neuralff/ops/__init__.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | from .fluid_ops import * 23 | from .grid_ops import * 24 | from .vector_ops import * 25 | from .physics_ops import * 26 | 27 | -------------------------------------------------------------------------------- /neuralff/ops/fluid_ops.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import torch 23 | from .grid_ops import sample_from_grid 24 | 25 | def semi_lagrangian_backtrace(coords, grid, velocity_field, timestep): 26 | """Perform semi-Lagrangian backtracing at continuous coordinates. 27 | 28 | Warning: PyTorch follows rather stupid conventions, so the flow field is [w, h] but the grid is [h, w] 29 | 30 | Args: 31 | coords (torch.FloatTensor): continuous coordinates in normalized coords [-1, 1] of size [N, 2] 32 | grid (torch.FloatTensor): grid of features of size [H, W, F] 33 | velocity_field (neuralff.Field): Eulerian vector field. This can be any Field representation. 34 | timestep (float) : timestep of the simulation 35 | 36 | Returns 37 | (torch.FloatTensor): backtracked values of [N, F] 38 | """ 39 | 40 | # Velocities aren't necessarily on the same grid as the coords 41 | velocities_at_coords = velocity_field.sample(coords) 42 | 43 | # Equation 3.22 from Doyub's book 44 | samples = sample_from_grid(coords - timestep * velocities_at_coords, grid) 45 | return samples 46 | 47 | def semi_lagrangian_advection(coords, grid, velocity_field, timestep): 48 | """Performs advection and updates the grid. 49 | 50 | This method is similar to the `semi_lagrangian_backtrace` function, but assumes the `coords` are 51 | perfectly aligned with the `grid`. 52 | 53 | Args: 54 | coords (torch.FloatTensor): continuous coordinates in normalized coords [-1, 1] of size [H, W, 2] 55 | grid (torch.FloatTensor): grid of features of size [H, W, F] 56 | velocity_field (neuralff.Field): Eulerian vector field. This can be any Field representation. 57 | timestep (float) : timestep of the simulation 58 | 59 | Returns 60 | (torch.FloatTensor): advaceted grid of [H, W, F] 61 | """ 62 | H, W = coords.shape[:2] 63 | samples = semi_lagrangian_backtrace(coords.reshape(-1, 2), grid, velocity_field, timestep) 64 | return samples.reshape(H, W, -1) 65 | 66 | 67 | -------------------------------------------------------------------------------- /neuralff/ops/grid_ops.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import torch 23 | import torch.nn.functional as F 24 | 25 | def normalized_grid_coords(height, width, aspect=True, device="cuda"): 26 | """Return the normalized [-1, 1] grid coordinates given height and width. 27 | 28 | Args: 29 | height (int) : height of the grid. 30 | width (int) : width of the grid. 31 | aspect (bool) : if True, use the aspect ratio to scale the coordinates, in which case the 32 | coords will not be normalzied to [-1, 1]. (Default: True) 33 | device : the device the tensors will be created on. 34 | """ 35 | aspect_ratio = width/height if aspect else 1.0 36 | 37 | window_x = torch.linspace(-1, 1, steps=width, device=device) * aspect_ratio 38 | window_y = torch.linspace(1, -1, steps=height, device=device) 39 | coord = torch.stack(torch.meshgrid(window_x, window_y, indexing='ij')).permute(2,1,0) 40 | return coord 41 | 42 | def sample_from_grid(coords, grid, padding_mode="reflection"): 43 | """Sample from a discrete grid at continuous coordinates. 44 | 45 | Args: 46 | coords (torch.FloatTensor): continuous coordinates in normalized coords [-1, 1] of size [N, 2] 47 | grid (torch.FloatTensor): grid of size [H, W, F]. 48 | 49 | Returns: 50 | (torch.FloatTensor): interpolated values of [N, F] 51 | 52 | """ 53 | N = coords.shape[0] 54 | sample_coords = coords.reshape(1, N, 1, 2) 55 | sample_coords = (sample_coords + 1.0) % 2.0 - 1.0 56 | samples = F.grid_sample(grid[None].permute(0,3,1,2), sample_coords, align_corners=True, 57 | padding_mode=padding_mode)[0,:,:,0].transpose(0,1) 58 | return samples 59 | 60 | -------------------------------------------------------------------------------- /neuralff/ops/physics_ops.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import torch 23 | from .vector_ops import divergence, laplacian, gradient 24 | 25 | def time_derivative( 26 | coords, velocity_field, timestep, initial_velocity=None): 27 | """Computes the time derivative (du/dt). 28 | 29 | Args: 30 | coords (torch.Tensor) : coordinates of shape [N, D] 31 | velocity_field (neuralff.Field) : velocity field function (aka u) 32 | timestep (float) : delta t of the simulation 33 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 34 | This is useful for preconditioning the velocity field. 35 | 36 | Returns: 37 | (torch.Tensor, torch.Tensor) : 38 | - time derivative of shape [N, D] 39 | - velocity of shape [N, D] 40 | """ 41 | u = velocity_field.sample(coords) 42 | if initial_velocity is None: 43 | dudt = (u - u.detach()) / timestep 44 | else: 45 | dudt = (u - initial_velocity) / timestep 46 | return dudt, u 47 | 48 | def material_derivative( 49 | coords, velocity_field, timestep, eps=1e-3, initial_velocity=None): 50 | """Computes the material derivative. 51 | 52 | We use the following definition of the material derivative (from Bridson): 53 | 54 | eq 1 : du/dt + (div u) u 55 | 56 | Args: 57 | coords (torch.Tensor) : coordinates of shape [N, D] 58 | velocity_field (neuralff.Field) : velocity field function (aka u) 59 | timestep (float) : delta t of the simulation 60 | eps (float) : the "grid spacing" for finite diff 61 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 62 | This is useful for preconditioning the velocity field. 63 | 64 | Returns: 65 | (torch.Tensor, torch.Tensor) : 66 | - material derivative of shape [N, 2] 67 | - divergence of shape [N, 1] 68 | """ 69 | dudt, u = time_derivative(coords, velocity_field, timestep, initial_velocity=initial_velocity) 70 | div_u = divergence(coords, velocity_field, eps=eps) 71 | return dudt + div_u * u, div_u 72 | 73 | def gravity(coords, timestep): 74 | """Returns the gravity vector constraint. 75 | 76 | Args: 77 | coords (torch.Tensor) : coordinates of shape [N, D] 78 | timestep (float) : delta t of the simulation 79 | 80 | Returns: 81 | (torch.Tensor) : gravity tensor of shape [N, D] 82 | """ 83 | g = torch.zeros_like(coords) 84 | g[..., 1] = 9.8 * 1e-4 85 | return g 86 | 87 | def divergence_free_loss(coords, velocity_field, eps=1e-3): 88 | """Computes the divergence-free equation loss. 89 | 90 | eq 1 : div u = 0 91 | 92 | Args: 93 | coords (torch.Tensor) : coordinates of shape [N, D] 94 | velocity_field (neuralff.Field) : velocity field function (aka u) 95 | eps (float) : the "grid spacing" for finite diff 96 | 97 | Returns: 98 | (torch.Tensor) : per-point loss of the divergence-free term of shape [N, 1] 99 | """ 100 | div_u = divergence(coords, velocity_field, eps=eps) 101 | return (div_u**2.0).sum(-1, keepdim=True) 102 | 103 | def advection_loss(coords, velocity_field, timestep, eps=1e-3, initial_velocity=None): 104 | """Computes the advection equation loss. 105 | 106 | We use the following definition (from Bridson) 107 | 108 | eq 1 : du/dt + (div u) u = 0 109 | 110 | Args: 111 | coords (torch.Tensor) : coordinates of shape [N, D] 112 | velocity_field (neuralff.Field) : velocity field function (aka u) 113 | timestep (float) : delta t of the simulation 114 | eps (float) : the "grid spacing" for finite diff 115 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 116 | This is useful for preconditioning the velocity field. 117 | 118 | Returns: 119 | (torch.Tensor) : per-point loss of the advection equation of shape [N, 1] 120 | """ 121 | mat_d, div_u = material_derivative(coords, velocity_field, timestep, eps=eps, initial_velocity=initial_velocity) 122 | return (mat_d**2.0).sum(-1, keepdim=True) 123 | 124 | def body_forces_loss(coords, velocity_field, timestep, eps=1e-3, initial_velocity=None): 125 | """Computes the body-forces equation loss. 126 | 127 | We use the following definition (from Bridson) 128 | 129 | eq 1 : du/dt = g 130 | 131 | Args: 132 | coords (torch.Tensor) : coordinates of shape [N, D] 133 | velocity_field (neuralff.Field) : velocity field function (aka u) 134 | timestep (float) : delta t of the simulation 135 | eps (float) : the "grid spacing" for finite diff 136 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 137 | This is useful for preconditioning the velocity field. 138 | 139 | Returns: 140 | (torch.Tensor) : per-point loss of the body forces equation of shape [N, 1] 141 | """ 142 | dudt, u = time_derivative(coords, velocity_field, timestep, initial_velocity=initial_velocity) 143 | g = gravity(coords, timestep) 144 | return ((g - dudt)**2.0).sum(-1, keepdim=True) 145 | 146 | def incompressibility_loss(coords, velocity_field, pressure_field, rho_field, timestep, eps=1e-3, initial_velocity=None): 147 | """Computes the incompressibility equation loss. 148 | 149 | We use the following definition (from Bridson) 150 | 151 | eq 1 : du/dt = (1/rho) grad p 152 | eq 2 : div u = 0 153 | 154 | Args: 155 | coords (torch.Tensor) : coordinates of shape [N, D] 156 | pressure_field (neuralff.Field) : pressure field function (aka p) 157 | rho_field (neuralff.Field) : the fluid density field (kg/m^d) 158 | timestep (float) : delta t of the simulation 159 | eps (float) : the "grid spacing" for finite diff 160 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 161 | This is useful for preconditioning the velocity field. 162 | 163 | Returns: 164 | (torch.Tensor) : per-point loss of the incompressibility equation of shape [N, 1] 165 | """ 166 | dudt, u = time_derivative(coords, velocity_field, timestep, initial_velocity=initial_velocity) 167 | div_u = divergence(coords, velocity_field, eps=eps) 168 | rho = rho_field.sample(coords) 169 | grad_p = gradient(coords, pressure_field, eps=eps) 170 | return ((rho * dudt - grad_p)**2.0).sum(-1, keepdim=True) + 100.0 * (div_u**2.0).sum(-1, keepdim=True) 171 | 172 | def euler_loss( 173 | coords, velocity_field, pressure_field, rho_field, 174 | timestep, eps=1e-6, initial_velocity=None): 175 | """Computes the incompressible Euler equation loss. 176 | 177 | We use the following definition of the Euler equation (from Bridson): 178 | 179 | eq 1 : du/dt + (div u) u + (1/rho) grad p = g 180 | eq 2 : div u = 0 181 | 182 | That is, this is the Navier-Stokes equation without the viscosity term. 183 | 184 | Args: 185 | coords (torch.Tensor) : coordinates to enforce Navier-Stokes of shape [N, D] 186 | velocity_field (neuralff.Field) : velocity field function (aka u) 187 | pressure_field (neuralff.Field) : pressure field function (aka p) 188 | rho_field (neuralff.Field) : the fluid density field (kg/m^d) 189 | timestep (float) : delta t of the simulation 190 | eps (float) : the "grid spacing" for finite diff 191 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 192 | This is useful for preconditioning the velocity field. 193 | 194 | Returns: 195 | (torch.Tensor) : per-point loss of the Euler equation of shape [N, 1] 196 | """ 197 | mat_d, div_u = material_derivative(coords, velocity_field, timestep, eps=eps, initial_velocity=initial_velocity) 198 | 199 | rho = rho_field.sample(coords) 200 | #grad_p = (1.0/rho) * gradient(coords, pressure_field, eps=eps) 201 | grad_p = gradient(coords, pressure_field, eps=eps) 202 | 203 | g = gravity(coords, timestep) 204 | 205 | #momentum_term = ((g - mat_d - grad_p)**2).sum(-1, keepdim=True) 206 | momentum_term = ((rho * mat_d + grad_p - rho * g)**2).sum(-1, keepdim=True) 207 | divergence_term = (div_u**2).sum(-1, keepdim=True) 208 | return momentum_term + 100.0 * divergence_term 209 | 210 | def navier_stokes_loss( 211 | coords, velocity_field, pressure_field, rho_field, 212 | timestep, nu=1e-5, eps=1e-3, initial_velocity=None): 213 | """Computes the Navier Stokes equation loss. 214 | 215 | We use the following definition of the Navier-Stokes equation (from Bridson): 216 | 217 | eq 1 : du/dt + (div u) u + (1/rho) grad p = g + nu (Lap u) 218 | eq 2 : div u = 0 219 | 220 | Args: 221 | coords (torch.Tensor) : coordinates to enforce Navier-Stokes of shape [N, D] 222 | velocity_field (neuralff.Field) : velocity field function (aka u) 223 | pressure_field (neuralff.Field) : pressure field function (aka p) 224 | rho_field (neuralff.Field) : the fluid density field (kg/m^d) 225 | timestep (float) : delta t of the simulation 226 | nu (float) : kinematic viscosity 227 | eps (float) : the "grid spacing" for finite diff 228 | initial_velocity (torch.Tensor) : if provided, will use this instead of sampling the initial velocity. 229 | This is useful for preconditioning the velocity field. 230 | 231 | Returns: 232 | (torch.Tensor) : per-point loss of the Navier-Stokes equation of shape [N, 1] 233 | """ 234 | mat_d, div_u = material_derivative(coords, velocity_field, timestep, eps=eps, initial_velocity=initial_velocity) 235 | 236 | rho = rho_field.sample(coords) 237 | grad_p = (1.0/rho) * gradient(coords, pressure_field, eps=eps) 238 | 239 | g = gravity(coords, timestep) 240 | 241 | diff = nu * laplacian(coords, velocity_field, eps=eps) 242 | 243 | momentum_term = ((g + diff - mat_d - grad_p)**2).sum(-1, keepdim=True) 244 | divergence_term = (div_u**2).sum(-1, keepdim=True) 245 | return momentum_term + divergence_term 246 | 247 | -------------------------------------------------------------------------------- /neuralff/ops/vector_ops.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import torch 23 | 24 | def gradient(u, f, method='finitediff', eps=1e-4): 25 | """Compute the gradient of the scalar field `f` with respect to the input `u`. 26 | 27 | Args: 28 | u (torch.Tensor) : the input to the function f of shape [..., D] 29 | f (function) : some scalar field (must support autodiff if using method='autodiff') 30 | method (str) : method for calculating the gradient. 31 | options: ['autodiff', 'finitediff'] 32 | 33 | Returns: 34 | (torch.Tensor) : gradient of shape [..., D] 35 | 36 | Finite diff currently assumes that `u` represents a 2D vector field (D=2). 37 | 38 | """ 39 | if method == 'autodiff': 40 | with torch.enable_grad(): 41 | u = u.requires_grad_(True) 42 | v = f(u) 43 | grad = torch.autograd.grad(v, u, 44 | grad_outputs=torch.ones_like(v), create_graph=True)[0] 45 | elif method == 'finitediff': 46 | assert(u.shape[-1] == 2 and "Finitediff only supports 2D vector fields") 47 | eps_x = torch.tensor([eps, 0.0], device=u.device) 48 | eps_y = torch.tensor([0.0, eps], device=u.device) 49 | 50 | grad = torch.cat([f(u + eps_x) - f(u - eps_x), 51 | f(u + eps_y) - f(u - eps_y)], dim=-1) 52 | grad = grad / (eps*2.0) 53 | else: 54 | raise NotImplementedError 55 | 56 | return grad 57 | 58 | def jacobian(u, f, method='finitediff', eps=1e-4): 59 | """Compute the Jacobian of the vector field `f` with respect to the input `u`. 60 | 61 | Args: 62 | u (torch.Tensor) : the input to the function f of shape [..., D] 63 | f (function) : some vector field (must support autodiff if using method='autodiff') with output dim [F] 64 | method (str) : method for calculating the Jacobian. 65 | options: ['autodiff', 'finitediff'] 66 | 67 | Returns: 68 | (torch.Tensor) : Jacobian of shape [..., F, D] 69 | 70 | Finite diff currently assumes that `u` represents a 2D vector field. 71 | """ 72 | if method == 'autodiff': 73 | raise NotImplementedError 74 | # The behaviour here is a bit mysterious to me... 75 | with torch.enable_grad(): 76 | j = torch.autograd.functional.jacobian(f, u, create_graph=True) 77 | elif method == 'finitediff': 78 | assert(u.shape[-1] == 2 and "Finitediff only supports 2D vector fields") 79 | eps_x = torch.tensor([eps, 0.0], device=u.device) 80 | eps_y = torch.tensor([0.0, eps], device=u.device) 81 | 82 | dfux = (f(u + eps_x) - f(u - eps_x))[..., None] 83 | dfuy = (f(u + eps_y) - f(u - eps_y))[..., None] 84 | 85 | # Check that the dims are ordered correctly 86 | return torch.cat([dfux, dfuy], dim=-1) / (eps*2.0) 87 | else: 88 | raise NotImplementedError 89 | return j 90 | 91 | def divergence(u, f, method='finitediff', eps=1e-4): 92 | """Compute the divergence of the vector field `f` with respect to the input `u`. 93 | 94 | Args: 95 | u (torch.Tensor) : the input to the function f of shape [..., D] 96 | f (function) : some vector field (must support autodiff if using method='autodiff') with output dim [D] 97 | method (str) : method for calculating the Jacobian. 98 | options: ['autodiff', 'finitediff'] 99 | 100 | Finite diff currently assumes that `u` represents a 2D vector field. 101 | 102 | Returns: 103 | (torch.Tensor) : divergence of shape [..., 1] 104 | """ 105 | if method == 'autodiff': 106 | raise NotImplementedError 107 | j = jacobian(u, f, method=method, eps=eps) 108 | return j.diagonal(offset=0, dim1=-1, dim2=-2).sum(-1) # Make sure it returns the correct diagonal 109 | if method == 'finitediff': 110 | eps_x = torch.tensor([eps, 0.0], device=u.device) 111 | eps_y = torch.tensor([0.0, eps], device=u.device) 112 | dfxux = f(u + eps_x)[...,0:1] - f(u - eps_x)[...,0:1] 113 | dfyuy = f(u + eps_y)[...,1:2] - f(u - eps_y)[...,1:2] 114 | return (dfxux + dfyuy)/(eps*2.0) 115 | 116 | def curl(u, f, method='finitediff', eps=1e-4): 117 | """Compute the curl of the vector field `f` with respect to the input `u`. 118 | 119 | Args: 120 | u (torch.Tensor) : the input to the function f of shape [..., D] 121 | f (function) : some vector field (must support autodiff if using method='autodiff') with output dim [D] 122 | method (str) : method for calculating the curl. 123 | options: ['autodiff', 'finitediff'] 124 | 125 | Finite diff currently assumes that `u` represents a 2D vector field. 126 | 127 | Returns: 128 | (torch.Tensor) : curl of shape [..., 1] 129 | """ 130 | if method == 'autodiff': 131 | raise NotImplementedError 132 | if method == 'finitediff': 133 | eps_x = torch.tensor([eps, 0.0], device=u.device) 134 | eps_y = torch.tensor([0.0, eps], device=u.device) 135 | dfyux = f(u + eps_x)[...,1:2] - f(u - eps_x)[...,1:2] 136 | dfxuy = f(u + eps_y)[...,0:1] - f(u - eps_y)[...,0:1] 137 | return (dfyux+dfxuy)/(eps*2.0) 138 | 139 | def laplacian(u, f, method='finitediff', eps=1e-4): 140 | """Compute the Laplacian of the vector field `f` with respect to the input `u`. 141 | 142 | Note: the Laplacian of a vector field is just the vector of Laplacians of its components. 143 | 144 | Args: 145 | u (torch.Tensor) : the input to the function f of shape [..., D] 146 | f (function) : some vector field (must support autodiff if using method='autodiff') with output dim [D] 147 | method (str) : method for calculating the Laplacian. 148 | options: ['autodiff', 'finitediff'] 149 | 150 | Finite diff currently assumes that `u` represents a 2D vector field. 151 | 152 | Returns: 153 | (torch.Tensor) : Laplacian of shape [..., 1] 154 | """ 155 | if method == 'autodiff': 156 | raise NotImplementedError 157 | if method == 'finitediff': 158 | fu = 2.0 * f(u) 159 | eps_x = torch.tensor([eps, 0.0], device=u.device) 160 | eps_y = torch.tensor([0.0, eps], device=u.device) 161 | dfux = f(u + eps_x) - fu + f(u - eps_x) 162 | dfuy = f(u + eps_y) - fu + f(u - eps_y) 163 | return (dfux + dfuy) / (eps**2.0) 164 | 165 | 166 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | # 3 | # Copyright (c) 2021 4 | # 5 | # Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | # this software and associated documentation files (the "Software"), to deal in 7 | # the Software without restriction, including without limitation the rights to 8 | # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 9 | # the Software, and to permit persons to whom the Software is furnished to do so, 10 | # subject to the following conditions: 11 | # 12 | # The above copyright notice and this permission notice shall be included in all 13 | # copies or substantial portions of the Software. 14 | # 15 | # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 17 | # FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 18 | # COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 19 | # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 20 | # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 21 | 22 | import os 23 | import sys 24 | from setuptools import setup, find_packages, dist 25 | import glob 26 | import logging 27 | 28 | import torch 29 | from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension 30 | 31 | PACKAGE_NAME = 'neuralff' 32 | DESCRIPTION = 'research on neural fluid fields' 33 | URL = 'https://github.com/tovacinni/neural-fluid-fields' 34 | AUTHOR = 'Towaki Takikawa, Mohammad Mozaffari' 35 | LICENSE = 'MIT License' 36 | version = '0.1.0' 37 | 38 | def get_extensions(): 39 | extra_compile_args = {'cxx': ['-O3']} 40 | define_macros = [] 41 | include_dirs = [] 42 | extensions = [] 43 | sources = glob.glob('neuralff/csrc/**/*.cpp', recursive=True) 44 | 45 | if len(sources) == 0: 46 | print("No source files found for extension, skipping extension compilation") 47 | return None 48 | 49 | if torch.cuda.is_available() or os.getenv('FORCE_CUDA', '0') == '1': 50 | define_macros += [("WITH_CUDA", None), ("THRUST_IGNORE_CUB_VERSION_CHECK", None)] 51 | sources += glob.glob('neuralff/csrc/**/*.cu', recursive=True) 52 | extension = CUDAExtension 53 | extra_compile_args.update({'nvcc': ['-O3']}) 54 | #include_dirs = get_include_dirs() 55 | else: 56 | assert(False, "CUDA is not available. Set FORCE_CUDA=1 for Docker builds") 57 | 58 | extensions.append( 59 | extension( 60 | name='neuralff._C', 61 | sources=sources, 62 | define_macros=define_macros, 63 | extra_compile_args=extra_compile_args, 64 | #include_dirs=include_dirs 65 | ) 66 | ) 67 | 68 | for ext in extensions: 69 | ext.libraries = ['cudart_static' if x == 'cudart' else x 70 | for x in ext.libraries] 71 | 72 | return extensions 73 | 74 | if __name__ == '__main__': 75 | setup( 76 | # Metadata 77 | name=PACKAGE_NAME, 78 | version=version, 79 | author=AUTHOR, 80 | description=DESCRIPTION, 81 | url=URL, 82 | license=LICENSE, 83 | python_requires='~=3.8', 84 | 85 | # Package info 86 | packages=['neuralff'], 87 | include_package_data=True, 88 | zip_safe=True, 89 | ext_modules=get_extensions(), 90 | cmdclass={ 91 | 'build_ext': BuildExtension.with_options(no_python_abi_suffix=True) 92 | } 93 | 94 | ) 95 | --------------------------------------------------------------------------------