├── .gitignore ├── README.md ├── figures ├── Animation_rho001.gif └── GrowthRate.png ├── pytorch ├── .gitignore ├── __init__.py ├── cylinder.py ├── fluid_net_simulate_one.py ├── fluid_net_train.py ├── lib │ ├── .gitignore │ ├── __init__.py │ ├── argument_parser.py │ ├── dataset_load.py │ ├── fluid │ │ ├── .gitignore │ │ ├── __init__.py │ │ ├── cell_type.py │ │ ├── cpp │ │ │ ├── .gitignore │ │ │ ├── __init__.py │ │ │ ├── advect_type.cpp │ │ │ ├── advect_type.h │ │ │ ├── advection.py │ │ │ ├── calc_line_trace.cpp │ │ │ ├── calc_line_trace.h │ │ │ ├── cell_type.h │ │ │ ├── fluids_init.cpp │ │ │ ├── fluids_init.h │ │ │ ├── grid.cpp │ │ │ ├── grid.h │ │ │ ├── jit.py │ │ │ ├── setup.py │ │ │ └── solve_linear_sys.py │ │ ├── flags_to_occupancy.py │ │ ├── geometry_utils.py │ │ ├── grid.py │ │ ├── init_conditions.py │ │ ├── set_wall_bcs.py │ │ ├── set_wall_bcs_inflow.py │ │ ├── set_wall_bcs_stick.py │ │ ├── source_terms.py │ │ ├── source_terms_test.py │ │ ├── util.py │ │ ├── velocity_divergence.py │ │ ├── velocity_update.py │ │ └── viscosity.py │ ├── load_manta_data.py │ ├── model.py │ ├── multi_scale_net.py │ ├── plot_field.py │ ├── plot_field_temp.py │ ├── simulate.py │ └── util_print.py ├── plot_5loss.py ├── plot_loss.py ├── plume.py ├── plumeConfig.yaml ├── print_output.py ├── rayleighTaylor.py ├── rayleighTaylorConfig.yaml └── trainConfig.yaml ├── solver_cpp ├── CMakeLists.txt ├── simulate │ ├── CMakeLists.txt │ └── simulate.cpp ├── src │ ├── CMakeLists.txt │ ├── advection │ │ ├── .gitignore │ │ ├── advect_type.cpp │ │ ├── advect_type.h │ │ ├── advection.cpp │ │ ├── advection.h │ │ ├── calc_line_trace.cpp │ │ └── calc_line_trace.h │ ├── boundaryCondition │ │ ├── .gitignore │ │ ├── bcs.cpp │ │ └── bcs.h │ ├── fluid.h │ ├── fluidnet_implementation │ │ ├── CMakeLists.txt │ │ ├── advect_type.cpp │ │ ├── advect_type.h │ │ ├── calc_line_trace.cpp │ │ ├── cell_type.h │ │ ├── grid.cpp │ │ ├── grid.h │ │ ├── init.cpp │ │ ├── int3.h │ │ ├── quadrants.h │ │ ├── stack_trace.cpp │ │ ├── tfluids.cpp │ │ ├── vec3.cpp │ │ └── vec3.h │ ├── grid │ │ ├── .gitignore │ │ ├── bool_conversion.cpp │ │ ├── bool_conversion.h │ │ ├── cell_type.h │ │ ├── grid.cpp │ │ └── grid.h │ ├── projection │ │ ├── .gitignore │ │ ├── div.cpp │ │ ├── div.h │ │ ├── solve_linear_sys.cpp │ │ ├── solve_linear_sys.h │ │ ├── update_vel.cpp │ │ └── update_vel.h │ └── sourceTerms │ │ ├── .gitignore │ │ ├── source_term.cpp │ │ └── source_term.h └── test │ ├── .gitignore │ ├── CMakeLists.txt │ ├── load_manta_data.h │ ├── plot_utils.h │ ├── test_fluid.cpp │ └── type_test.h └── trained_models └── ScaleNet_ShortTerm_LongTermLoss ├── ScaleNet_ShortTerm_LongTermLoss_saved.py ├── convModel_conf.pth ├── convModel_lastEpoch.pth ├── convModel_lastEpoch_best.pth └── convModel_mconf.pth /.gitignore: -------------------------------------------------------------------------------- 1 | test_data/ 2 | build*/ 3 | /*data* 4 | !*/data* 5 | src 6 | -------------------------------------------------------------------------------- /figures/Animation_rho001.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/figures/Animation_rho001.gif -------------------------------------------------------------------------------- /figures/GrowthRate.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/figures/GrowthRate.png -------------------------------------------------------------------------------- /pytorch/.gitignore: -------------------------------------------------------------------------------- 1 | misc 2 | figures 3 | plot 4 | model 5 | data 6 | data2 7 | __pycache__ 8 | pybind11 9 | *.pyc 10 | -------------------------------------------------------------------------------- /pytorch/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/pytorch/__init__.py -------------------------------------------------------------------------------- /pytorch/lib/.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | *resume.py 3 | *model_print.py 4 | -------------------------------------------------------------------------------- /pytorch/lib/__init__.py: -------------------------------------------------------------------------------- 1 | from .dataset_load import FluidNetDataset 2 | from .util_print import summary 3 | from .multi_scale_net import MultiScaleNet 4 | from .model import FluidNet 5 | from .simulate import simulate 6 | from .plot_field import plotField 7 | from .argument_parser import SmartFormatter 8 | -------------------------------------------------------------------------------- /pytorch/lib/argument_parser.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | class SmartFormatter(argparse.HelpFormatter): 4 | 5 | def _split_lines(self, text, width): 6 | if text.startswith('R|'): 7 | return text[2:].splitlines() 8 | # this is the RawTextHelpFormatter._split_lines 9 | return argparse.HelpFormatter._split_lines(self, text, width) 10 | -------------------------------------------------------------------------------- /pytorch/lib/dataset_load.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.utils.data import Dataset, DataLoader 3 | import torch.multiprocessing as mp 4 | import glob 5 | from .load_manta_data import loadMantaFile 6 | 7 | class FluidNetDataset(Dataset): 8 | """Fluid Net dataset.""" 9 | 10 | def __init__(self, conf, prefix, save_dt, preprocess=False, resume=False, pr_n_threads=0): 11 | 12 | self.conf = conf.copy() 13 | self.mconf = self.conf['modelParam'] 14 | del self.conf['modelParam'] 15 | 16 | self.prefix = prefix 17 | self.save_dt = save_dt 18 | self.data_dir = conf['dataDir'] 19 | self.dataset = conf['dataset'] 20 | self.n_threads = pr_n_threads # Num of threads to preprocess from .bin to .pt 21 | 22 | self.base_dir = self._get_base_dir(conf['dataDir']) 23 | self.scenes_folders = sorted(glob.os.listdir(self.base_dir)) 24 | 25 | 26 | # Check number of scenes 27 | self.n_scenes = len(self.scenes_folders) 28 | self.scene_0 = int(self.scenes_folders[0]) 29 | 30 | # Check how many timesteps per scene there are. 31 | #self.step_per_scene = len(glob.glob(glob.os.path.join(self.base_dir, \ 32 | # self.scenes_folders[0], '*[0-9].bin'))) 33 | # TODO (Remove this): Hard-coded (there are 64 steps per scene) 34 | self.step_per_scene = 64 35 | 36 | self.pr_loader = loadMantaFile 37 | self.loader = torch.load 38 | 39 | # Log file (exists if dataset was preprocessed in the past) 40 | # It contains a dict with the name of data tensors and target tensors 41 | self.pr_log = {} 42 | self.logname = glob.os.path.join(self.data_dir, self.dataset, \ 43 | 'preprocessed_' + self.dataset + '_' + self.prefix + '.txt') 44 | 45 | # Pre-process 46 | if not resume: 47 | # Checking if pre-processing is needed: 48 | yes = {'yes','y', 'ye', ''} 49 | no = {'no','n'} 50 | if (glob.os.path.isfile(self.logname)): 51 | if preprocess: 52 | print('For dataset ' + str(self.base_dir)) 53 | print('a log file exists showing a preprocessing in the past.') 54 | self.pr_log = torch.load(self.logname) 55 | print(self.pr_log) 56 | print('Do you want to preprocess the dataset again? It takes some time. [y/n]') 57 | choice = input().lower() 58 | if choice in yes: 59 | self.preprocess() 60 | elif choice in no: 61 | # We check if file exits in __getitem__ 62 | pass 63 | else: 64 | sys.stdout.write("Please respond with 'yes' or 'no'") 65 | else: 66 | self.pr_log = torch.load(self.logname) 67 | pass 68 | else: 69 | print('No log file found for ' + str(self.base_dir) \ 70 | + ' dataset. Preprocessing automatically.') 71 | self.preprocess() 72 | 73 | else: 74 | if (glob.os.path.isfile(self.logname)): 75 | print() 76 | print('For dataset ' + str(self.base_dir)) 77 | print('a log file exists showing a preprocessing in the past.') 78 | print('OK to proceed with restarting') 79 | self.pr_log = torch.load(self.logname) 80 | else: 81 | print() 82 | print('No log file found, please create one by preprocessing the dataset. Set resume to false.') 83 | sys.exit() 84 | 85 | # Depending on inputs and loss, we will load different data in __getitem__ 86 | self.nx = self.pr_log['nx'] 87 | self.ny = self.pr_log['ny'] 88 | self.nz = self.pr_log['nz'] 89 | 90 | self.is3D = self.pr_log['is3D'] 91 | self.inputChan = self.mconf['inputChannels'] 92 | self.inDims = 1 # Flags is always passed 93 | 94 | if self.inputChan['div']: 95 | self.inDims += 1 96 | if self.inputChan['pDiv']: 97 | self.inDims += 1 98 | if self.inputChan['UDiv']: 99 | if self.is3D: 100 | self.inDims += 3 101 | else: 102 | self.inDims += 2 103 | 104 | self.mconf['inputDim'] = self.inDims 105 | self.mconf['is3D'] = self.is3D 106 | 107 | def createConfDict(self): 108 | return self.conf, self.mconf 109 | 110 | def preprocess(self): 111 | # Preprocess the dataset from .bin to .pt (much faster I/O) 112 | p = mp.Pool(self.n_threads) 113 | start = self.scene_0 * self.step_per_scene 114 | end = start + self.__len__() 115 | data_inputs = [idx for idx in range(start, end)] 116 | print('Pre-processing dataset:') 117 | try: 118 | p.map(self.__getitembin__, data_inputs) 119 | print('Pre-processing succeeded') 120 | self.pr_log = {'data': ['pDiv', 'UDiv', 'flagsDiv', 'densityDiv'], 121 | 'target' : ['p', 'U', 'density'], 'is3D' : False, 'nx' : 128, 122 | 'ny' : 128, 'nz' : 1} 123 | print('Log is now:') 124 | print(self.pr_log) 125 | torch.save(self.pr_log, self.logname) 126 | except: 127 | print('Pre-processing failed') 128 | 129 | def __len__(self): 130 | return self.n_scenes * self.step_per_scene 131 | 132 | # Used only in preprocessing 133 | def __getitembin__(self, idx): 134 | cur_scene = idx // self.step_per_scene 135 | cur_timestep = (idx % (self.step_per_scene)) * self.save_dt 136 | data_file = glob.os.path.join(self.base_dir, '{0:06d}'.format(cur_scene), \ 137 | '{0:06d}.bin'.format(cur_timestep)) 138 | data_div_file = glob.os.path.join(self.base_dir, '{0:06d}'.format(cur_scene), \ 139 | '{0:06d}_divergent.bin'.format(cur_timestep)) 140 | assert glob.os.path.isfile(data_file), 'Data file ' + data_file + ' does not exists' 141 | assert glob.os.path.isfile(data_div_file), 'Data file does not exists' 142 | p, U, flags, density, is3D = self.pr_loader(data_file) 143 | pDiv, UDiv, flagsDiv, densityDiv, is3DDiv = self.pr_loader(data_div_file) 144 | 145 | assert is3D == is3DDiv, '3D flag is inconsistent!' 146 | assert torch.equal(flags, flagsDiv), 'Flags are not equal for idx ' + str(idx) 147 | 148 | data = torch.cat([pDiv, UDiv, flagsDiv, densityDiv, p, U, density], 1) 149 | 150 | 151 | save_file = glob.os.path.join(self.base_dir, '{0:06d}'.format(cur_scene), \ 152 | '{0:06d}_pyTen.pt'.format(cur_timestep)) 153 | torch.save(data, save_file) 154 | 155 | 156 | # Actual data loader 157 | def __getitem__(self, idx): 158 | cur_scene = idx // self.step_per_scene 159 | cur_timestep = (idx % (self.step_per_scene)) * self.save_dt 160 | data_file = glob.os.path.join(self.base_dir, '{0:06d}'.format(cur_scene), \ 161 | '{0:06d}_pyTen.pt'.format(cur_timestep)) 162 | assert glob.os.path.isfile(data_file), 'Data file ' + data_file + ' does not exists' 163 | torch_file = self.loader(data_file) 164 | 165 | if (self.is3D): 166 | data = torch_file[0,0:6] 167 | target = torch_file[0,6:11] 168 | else: 169 | data = torch_file[0,0:5] 170 | target = torch_file[0,5:9] 171 | 172 | # # data indexes | | 173 | # # (dim 1) | 2D | 3D 174 | # # ---------------------------------------- 175 | # # DATA: 176 | # # pDiv | 0 | 0 177 | # # UDiv | 1:3 | 1:4 178 | # # flags | 3 | 4 179 | # # densityDiv | 4 | 5 180 | # # TARGET: 181 | # # p | 5 | 6 182 | # # U | 6:8 | 7:10 183 | # # density | 8 | 10 184 | 185 | return data, target 186 | 187 | 188 | def _get_base_dir(self, data_dir): 189 | return glob.os.path.join(data_dir, self.dataset, self.prefix) 190 | 191 | 192 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/__init__.py: -------------------------------------------------------------------------------- 1 | from .cell_type import CellType 2 | from .grid import getDx, getCentered 3 | from .set_wall_bcs import setWallBcs 4 | from .set_wall_bcs_stick import setWallBcsStick 5 | from .flags_to_occupancy import flagsToOccupancy 6 | from .velocity_divergence import velocityDivergence 7 | from .velocity_update import velocityUpdate 8 | from .source_terms import addBuoyancy, addGravity 9 | from .viscosity import addViscosity 10 | from .geometry_utils import createCylinder, createBox2D 11 | from .util import emptyDomain 12 | from .init_conditions import createPlumeBCs, createRayleighTaylorBCs 13 | from .cpp.advection import correctScalar, advectScalar, advectVelocity 14 | from .cpp.solve_linear_sys import solveLinearSystemJacobi 15 | 16 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cell_type.py: -------------------------------------------------------------------------------- 1 | from enum import IntEnum 2 | 3 | # We use the same convention as Mantaflow and FluidNet 4 | 5 | class CellType(IntEnum): 6 | TypeNone = 0 7 | TypeFluid = 1 8 | TypeObstacle = 2 9 | TypeEmpty = 4 10 | TypeInflow = 8 11 | TypeOutflow = 16 12 | TypeOpen = 32 13 | TypeStick = 128 14 | TypeReserved = 256 15 | 16 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/.gitignore: -------------------------------------------------------------------------------- 1 | build 2 | advection_cpp* 3 | dist 4 | *.txt 5 | __pycache__ 6 | fluidnet_cpp.egg-info 7 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/pytorch/lib/fluid/cpp/__init__.py -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/advect_type.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "torch/extension.h" 3 | #include "advect_type.h" 4 | 5 | AdvectMethod StringToAdvectMethod(const std::string& str) { 6 | if (str == "eulerFluidNet") { 7 | return ADVECT_EULER_FLUIDNET; 8 | } else if (str == "maccormackFluidNet") { 9 | return ADVECT_MACCORMACK_FLUIDNET; 10 | } else { 11 | std::stringstream ss; 12 | ss << "advection method (" << str << ") not supported (options " 13 | << "are: eulerFluidNet, maccormackFluidNet)"; 14 | AT_ERROR("Error: Advection method not supported"); 15 | } 16 | } 17 | 18 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/advect_type.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | typedef enum { 6 | ADVECT_EULER_FLUIDNET = 0, 7 | ADVECT_MACCORMACK_FLUIDNET = 1, 8 | } AdvectMethod; 9 | 10 | AdvectMethod StringToAdvectMethod(const std::string& str); 11 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/advection.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import fluidnet_cpp 3 | 4 | def __check_advection_method__(method): 5 | assert (method == 'eulerFluidNet' or method == 'maccormackFluidNet'), \ 6 | 'Error: Advection method not supported. Options are: \ 7 | maccormackFluidNet, eulerFluidNet' 8 | 9 | def correctScalar(dt, src, div, flags): 10 | maskFluid = flags.eq(1) 11 | src.masked_scatter_(maskFluid, \ 12 | (src + dt*0.5*src*div).masked_select(maskFluid)) 13 | 14 | def advectScalar(dt, src, U, flags, method = 'maccormackFluidNet', boundary_width = 1, 15 | sample_outside_fluid = False, maccormack_strength = 0.75): 16 | r"""Advects scalar field src by the input vel field U 17 | 18 | Arguments: 19 | dt (float): timestep in seconds. 20 | src (Tensor): input scalar field, to be advected. 21 | Shape is (batch,1,D,H,W) with D=1 for 2D 22 | and D>1 for 3D simulations. 23 | U (Tensor): input velocity field. 24 | Shape is (batch,2/3,D,H,W) with D=1 for 2D 25 | and D>1 for 3D simulations. 26 | flags (Tensor): Input occupancy grid. 27 | method (string, optional): Sets the method of advection. 28 | Options are eulerFluidNet and maccormackFluidNet. 29 | Defaults to maccormackFluidNet. 30 | boundary_width (int, optional): width of fluid domain boundary. 31 | Defaults to 1. 32 | sample_outside_fluid(bool, optional): For density advection, we do not want 33 | to advect values inside non-fluid cells and so this should be set to false. 34 | For other quantities (like temperature), this should be true. 35 | Defaults to ''False''. 36 | maccormack_strength (float, optional): A strength parameter that 37 | will make the advection eularian (with values interpolating in between). A 38 | value of 1 (which implements the update from An Unconditionally Stable 39 | MaCormack Method) tends to add too much high-frequency detail. 40 | """ 41 | #Check sizes 42 | __check_advection_method__(method) 43 | 44 | assert src.dim() == 5 and U.dim() == 5 and flags.dim() == 5, "Dimension mismatch" 45 | assert flags.size(1) == 1, "flags is not scalar" 46 | 47 | bsz = flags.size(0) 48 | d = flags.size(2) 49 | h = flags.size(3) 50 | w = flags.size(4) 51 | 52 | is3D = U.size(1) == 3 53 | if (not is3D): 54 | assert d == 1, "2D velocity field but zdepth > 1" 55 | assert U.size(1) == 2, "2D velocity field must have only 2 channels" 56 | 57 | # TODO: Debug 3D 58 | assert is3D == False, '3D is not supported yet!' 59 | assert U.size(0) == bsz and U.size(2) == d and \ 60 | U.size(3) == h and U.size(4) == w, "Size mismatch" 61 | assert U.is_contiguous() and flags.is_contiguous() and \ 62 | src.is_contiguous(), "Input is not contiguous" 63 | 64 | s_dst = fluidnet_cpp.advect_scalar(dt, src, U, flags, method, 65 | boundary_width, sample_outside_fluid, maccormack_strength) 66 | return s_dst 67 | 68 | def advectVelocity(dt, orig, U, flags, method = 'maccormackFluidNet', boundary_width = 1, 69 | maccormack_strength = 0.75): 70 | r"""Advects velocity field orig by velocity field U 71 | 72 | Arguments: 73 | dt (float): timestep in seconds. 74 | orig (Tensor): velocity field to be advected. 75 | Shape is (batch,2/3,D,H,W) with D=1 for 2D 76 | and D>1 for 3D simulations. 77 | U (Tensor): non-divergent velocity field to advect orig. 78 | Shape is (batch,2/3,D,H,W) with D=1 for 2D 79 | and D>1 for 3D simulations. 80 | flags (Tensor): Input occupancy grid. 81 | method (string, optional): Sets the method of advection. 82 | Options are eulerFluidNet and maccormackFluidNet. 83 | Defaults to maccormackFluidNet. 84 | boundary_width (int, optional): width of fluid domain boundary. 85 | Defaults to 1. 86 | maccormack_strength (float, optional): A strength parameter that 87 | will make the advection eularian (with values interpolating in between). A 88 | value of 1 (which implements the update from An Unconditionally Stable 89 | MaCormack Method) tends to add too much high-frequency detail. 90 | """ 91 | 92 | #Check sizes 93 | assert U.dim() == 5 and orig.dim() == 5 and flags.dim() == 5, "Dimension mismatch" 94 | assert flags.size(1) == 1, "flags is not scalar" 95 | 96 | bsz = flags.size(0) 97 | d = flags.size(2) 98 | h = flags.size(3) 99 | w = flags.size(4) 100 | 101 | is3D = U.size(1) == 3 102 | if (not is3D): 103 | assert d == 1, "2D velocity field but zdepth > 1" 104 | assert orig.size(1) == 2, "2D velocity field must have only 2 channels" 105 | assert U.size(1) == 2, "2D velocity field must have only 2 channels" 106 | 107 | # TODO: Debug 3D 108 | assert is3D == False, '3D is not supported yet!' 109 | assert U.size(0) == bsz and U.size(2) == d and \ 110 | U.size(3) == h and U.size(4) == w, "Size mismatch" 111 | assert orig.size(0) == bsz and orig.size(2) == d and \ 112 | orig.size(3) == h and orig.size(4) == w, "Size mismatch" 113 | assert U.is_contiguous() and orig.is_contiguous() and flags.is_contiguous(), "Input is not contiguous" 114 | 115 | U_dst = fluidnet_cpp.advect_vel(dt, orig, U, flags, method, 116 | boundary_width, maccormack_strength) 117 | 118 | return U_dst 119 | 120 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/calc_line_trace.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "torch/extension.h" 4 | #include "cell_type.h" 5 | 6 | // See the .cpp file for detailed comments. 7 | namespace fluid { 8 | 9 | typedef at::Tensor T; 10 | 11 | void getPixelCenter(const T& pos, T& ix); 12 | 13 | T isOutOfDomain(const T& pos, const T& flags); 14 | 15 | T isBlockedCell(const T& pos, const T& flags); 16 | 17 | void clampToDomain(T& pos, const T& flags); 18 | 19 | typedef enum Quadrants { 20 | RIGHT = 0, 21 | LEFT = 1, 22 | MIDDLE = 2 23 | } Quadrants; 24 | 25 | T HitBoundingBox(const T& minB, const T& maxB, 26 | const T& origin, const T& dir, 27 | const T& mask, T& coord); 28 | 29 | T calcRayBoxIntersection(const T& pos, const T& dt, const T& ctr, 30 | const float hit_margin, const T& mask, T& ipos); 31 | 32 | T calcRayBorderIntersection(const T& pos, const T& next_pos, 33 | const T& flags, const float hit_margin, 34 | const T& mOutDom, 35 | T& ipos); 36 | 37 | void calcLineTrace(const T& pos, const T& delta, const T& flags, 38 | T& new_pos, const bool do_line_trace); 39 | 40 | } // namespace fluid 41 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/cell_type.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | namespace fluid { 4 | // These are the same enum values used in Manta. We can't include grid.h 5 | // from Manta without pulling in the entire library, so we'll just redefine 6 | // them here. 7 | enum CellType { 8 | TypeNone = 0, 9 | TypeFluid = 1, 10 | TypeObstacle = 2, 11 | TypeEmpty = 4, 12 | TypeInflow = 8, 13 | TypeOutflow = 16, 14 | TypeOpen = 32, 15 | TypeStick = 128, 16 | TypeReserved = 256, 17 | TypeZeroPressure = (1<<15) 18 | }; 19 | 20 | } // namespace fluid 21 | 22 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/fluids_init.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | #include "torch/extension.h" 7 | #include "grid.h" 8 | #include "cell_type.h" 9 | #include "advect_type.h" 10 | #include "calc_line_trace.h" 11 | 12 | namespace fluid { 13 | 14 | typedef at::Tensor T; 15 | 16 | T SemiLagrangeEulerFluidNet 17 | ( 18 | T& flags, T& vel, T& src, T& maskBorder, 19 | float dt, float order_space, 20 | T& i, T& j, T& k, 21 | const bool line_trace, 22 | const bool sample_outside_fluid 23 | ); 24 | 25 | T SemiLagrangeEulerFluidNetSavePos 26 | ( 27 | T& flags, T& vel, T& src, T& maskBorder, 28 | float dt, float order_space, 29 | T& i, T& j, T& k, 30 | const bool line_trace, 31 | const bool sample_outside_fluid, 32 | T& pos 33 | ); 34 | 35 | T MacCormackCorrect 36 | ( 37 | T& flags, const T& old, 38 | const T& fwd, const T& bwd, 39 | const float strength, 40 | bool is_levelset 41 | ); 42 | 43 | T getClampBounds 44 | ( 45 | const T& src, const T& pos, const T& flags, 46 | const bool sample_outside, 47 | T& clamp_min, T& clamp_max 48 | ); 49 | 50 | T MacCormackClampFluidNet( 51 | T& flags, T& vel, 52 | const T& dst, const T& src, 53 | const T& fwd, const T& fwd_pos, 54 | const T& bwd_pos, const bool sample_outside 55 | ); 56 | 57 | // Advect scalar field 'p' by the input vel field 'u'. 58 | // 59 | // @input dt - timestep (seconds). 60 | // @input s - input scalar field to advect 61 | // @input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 62 | // @input flags - input occupancy grid 63 | // @input sDst - Returned scalar field. 64 | // @input method - OPTIONAL - "eulerFluidNet", "maccormackFluidNet" 65 | // @param boundaryWidth - OPTIONAL - boundary width. (default 1) 66 | // @param sampleOutsideFluid - OPTIONAL - For density advection we do not want 67 | // to advect values inside non-fluid cells and so this should be set to false. 68 | // For other quantities (like temperature), this should be true. 69 | // @param maccormackStrength - OPTIONAL - (default 0.75) A strength parameter 70 | // will make the advection eularian (with values interpolating in between). A 71 | // value of 1 (which implements the update from An Unconditionally Stable 72 | // MaCormack Method) tends to add too much high-frequency detail 73 | T advectScalar 74 | ( 75 | float dt, T src, T U, T flags, 76 | const std::string method_str, 77 | int bnd, 78 | const bool sample_outside_fluid, 79 | const float maccormack_strength 80 | ); 81 | 82 | T SemiLagrangeEulerFluidNetMAC 83 | ( 84 | T& flags, T& vel, T& src, T& mask, 85 | float dt, float order_space, 86 | const bool line_trace, 87 | T& i, T& j, T& k 88 | ); 89 | 90 | T MacCormackCorrectMAC 91 | ( 92 | T& flags, const T& old, 93 | const T& fwd, const T& bwd, 94 | const float strength, 95 | T& i, T& j, T& k 96 | ); 97 | 98 | T doClampComponentMAC 99 | ( 100 | int chan, 101 | const T& flags, const T& dst, 102 | const T& orig, const T& fwd, 103 | const T& pos, const T& vel 104 | ); 105 | 106 | T MacCormackClampMAC 107 | ( 108 | const T& flags, const T& vel, const T& dval, 109 | const T& orig, const T& fwd, const T& mask, 110 | float dt, 111 | const T& i, const T& j, const T& k 112 | ); 113 | 114 | // Advect velocity field 'u' by itself and store in uDst. 115 | // 116 | // @input dt - timestep (seconds). 117 | // @input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 118 | // @input flags - input occupancy grid 119 | // @input UDst - Returned velocity field. 120 | // @input method - OPTIONAL - "eulerFluidNet", "maccormackFluidNet" (default) 121 | // @input boundaryWidth - OPTIONAL - boundary width. (default 1) 122 | // @input maccormackStrength - OPTIONAL - (default 0.75) A strength parameter 123 | // will make the advection more 1st order (with values interpolating in 124 | // between). A value of 1 (which implements the update from "An Unconditionally 125 | // Stable MaCormack Method") tends to add too much high-frequency detail. 126 | T advectVel 127 | ( 128 | float dt, T orig, T U, T flags, 129 | const std::string method_str, 130 | int bnd, 131 | const float maccormack_strength 132 | ); 133 | 134 | std::vector solveLinearSystemJacobi 135 | ( 136 | T flags, 137 | T div, 138 | const bool is_3d, 139 | const float p_tol, 140 | const int max_iter, 141 | const bool verbose 142 | ); 143 | 144 | } // namespace fluid 145 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/grid.h: -------------------------------------------------------------------------------- 1 | #include "torch/extension.h" 2 | 3 | namespace fluid { 4 | 5 | typedef at::Tensor T; 6 | 7 | float getDx(at::Tensor self); 8 | 9 | T interpol(const T& self, const T& pos); 10 | 11 | void interpol1DWithFluid( 12 | const T& val_a, const T& is_fluid_a, 13 | const T& val_b, const T& is_fluid_b, 14 | const T& t_a, const T& t_b, 15 | T& is_fluid_ab, T& val_ab); 16 | 17 | T interpolWithFluid(const T& self, const T& flags, const T& pos); 18 | 19 | T getCentered(const T& self); 20 | 21 | T getAtMACX(const T& self); 22 | T getAtMACY(const T& self); 23 | T getAtMACZ(const T& self); 24 | 25 | T interpolComponent(const T& self, const T& pos, int c); 26 | 27 | T curl(const T& self); 28 | 29 | } // namespace fluid 30 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/jit.py: -------------------------------------------------------------------------------- 1 | from torch.utils.cpp_extension import load 2 | fluidnet_cpp = load( 3 | name="fluidnet_cpp", 4 | sources=[ 5 | "grid.cpp", 6 | "advect_type.cpp", 7 | "calc_line_trace.cpp", 8 | "fluids_init.cpp" 9 | ], 10 | verbose=True) 11 | help(fluidnet_cpp) 12 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import CppExtension, BuildExtension 3 | 4 | setup( 5 | name='fluidnet_cpp', 6 | ext_modules=[ 7 | CppExtension( 8 | 'fluidnet_cpp', 9 | [ 10 | 'grid.cpp', 11 | 'advect_type.cpp', 12 | 'calc_line_trace.cpp', 13 | 'fluids_init.cpp' 14 | ]), 15 | ], 16 | cmdclass={ 17 | 'build_ext': BuildExtension 18 | }) 19 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/cpp/solve_linear_sys.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import fluidnet_cpp 3 | 4 | def solveLinearSystemJacobi(flags, div, is_3d=False, p_tol=1e-5, max_iter=1000, verbose=False): 5 | r"""Solves the linear system using the Jacobi method. 6 | Note: Since we don't receive a velocity field, we need to receive the is3D 7 | flag from the caller. 8 | 9 | Arguments: 10 | flags (Tensor): Input occupancy grid. 11 | div (Tensor): The velocity divergence. 12 | is_3d (Tensor, optional): If True, a 3D domain is expected. 13 | p_tol (float, optional): ||p - p_prev|| termination condition. 14 | Defaults 1e-5. 15 | max_iter (int, optional): Maximum number of Jacobi iterations. 16 | Defaults 1000. 17 | verbose (bool, optional). Defaults False. 18 | Output: 19 | p (Tensor): Pressure field 20 | p_tol: Maximum residual accross all batches. 21 | 22 | """ 23 | #Check sizes 24 | 25 | assert div.dim() == 5 and flags.dim() == 5, "Dimension mismatch" 26 | assert flags.size(1) == 1, "flags is not scalar" 27 | 28 | bsz = flags.size(0) 29 | d = flags.size(2) 30 | h = flags.size(3) 31 | w = flags.size(4) 32 | 33 | assert div.is_same_size(flags), "Size mismatch" 34 | 35 | assert flags.is_contiguous() and div.is_contiguous(), "Input is not contiguous" 36 | 37 | p, p_tol = fluidnet_cpp.solve_linear_system(flags, div, is_3d, \ 38 | p_tol, max_iter, verbose) 39 | 40 | return p, p_tol 41 | 42 | 43 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/flags_to_occupancy.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from . import CellType 3 | 4 | 5 | 6 | def flagsToOccupancy(flags): 7 | r"""Transforms the flags tensor to occupancy tensor (0 for fluids, 1 for obstacles). 8 | 9 | Arguments: 10 | flags (Tensor): Input occupancy grid. 11 | Output: 12 | occupancy (Tensor): Output occupancy grid (0s and 1s). 13 | """ 14 | occupancy = flags.clone() 15 | flagsFluid = occupancy.eq(CellType.TypeFluid) 16 | flagsObstacle = occupancy.eq(CellType.TypeObstacle) 17 | occupancy.masked_fill_(flagsFluid, 0) 18 | occupancy.masked_fill_(flagsObstacle, 1) 19 | return occupancy 20 | 21 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/geometry_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from . import CellType 3 | 4 | def createCylinder(batch_dict, centerX, centerY, radius): 5 | r"""Adds (inplace) a cylinder to the flags tensor in input batch. 6 | 7 | Arguments: 8 | batch_dict (dict): Input batch of tensor. 9 | centerX (float): X-coordinate of cylinder center. 10 | centerY (float): Y-coordinate of cylinder center. 11 | radius (float): Radius of cylinder. 12 | 13 | """ 14 | cuda = torch.device('cuda') 15 | assert 'flags' in batch_dict, 'Error: flags key is not in batch dict' 16 | flags = batch_dict['flags'] 17 | assert flags.dim() == 5, 'Input flags must have 5 dimensions' 18 | assert flags.size(0) == 1, 'Only batches of size 1 allowed (inference)' 19 | xdim = flags.size(4) 20 | ydim = flags.size(3) 21 | zdim = flags.size(2) 22 | is3D = (zdim > 1) 23 | 24 | # Create the cylinder 25 | X = torch.arange(0, xdim, device=cuda).view(xdim).expand((1,ydim,xdim)) 26 | Y = torch.arange(0, ydim, device=cuda).view(ydim, 1).expand((1,ydim,xdim)) 27 | 28 | dist_from_center = (X - centerX).pow(2) + (Y-centerY).pow(2) 29 | mask_cylinder = dist_from_center <= radius*radius 30 | 31 | flags.masked_fill_(mask_cylinder, CellType.TypeObstacle) 32 | batch_dict['flags'] = flags 33 | 34 | def createBox2D(batch_dict, x0, x1, y0, y1): 35 | r"""Adds (inplace) a 2D Box to the flags tensor in input batch. 36 | 37 | Arguments: 38 | batch_dict (dict): Input batch of tensor. 39 | x0 (float): bottom-left x-coordinate. 40 | y0 (float): bottom-left y-coordinate. 41 | x1 (float): upper-right x-coordinate. 42 | y1 (float): upper-right y-coordinate. 43 | 44 | """ 45 | cuda = torch.device('cuda') 46 | assert 'flags' in batch_dict, 'Error: flags key is not in batch dict' 47 | flags = batch_dict['flags'] 48 | assert flags.dim() == 5, 'Input flags must have 5 dimensions' 49 | assert flags.size(0) == 1, 'Only batches of size 1 allowed (inference)' 50 | xdim = flags.size(4) 51 | ydim = flags.size(3) 52 | zdim = flags.size(2) 53 | is3D = (zdim > 1) 54 | 55 | # Create the cylinder 56 | X = torch.arange(0, xdim, device=cuda).view(xdim).expand((1,ydim,xdim)) 57 | Y = torch.arange(0, ydim, device=cuda).view(ydim, 1).expand((1,ydim,xdim)) 58 | 59 | mask_box_2D = (X >= x0).__and__(X < x1).__and__\ 60 | (Y >= y1).__and__(Y < y1) 61 | 62 | flags.masked_fill_(mask_box2_2D, CellType.TypeObstacle) 63 | batch_dict['flags'] = flags 64 | 65 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/grid.py: -------------------------------------------------------------------------------- 1 | import torch 2 | # Some utils for printing to VTK Tensors 3 | def getDx(self): 4 | grid_size_max = max(max(self.size(2), self.size(3)), self.size(4)) 5 | return (1.0 / grid_size_max) 6 | 7 | def getCentered(self): 8 | 9 | cuda = torch.device('cuda') 10 | bsz = self.size(0) 11 | d = self.size(2) 12 | h = self.size(3) 13 | w = self.size(4) 14 | 15 | is3D = (d > 1) 16 | 17 | c_vel_x = torch.zeros_like(self[:,0]) 18 | c_vel_x[:,:,:,:-1] = 0.5 * (self[:,0,:,:,0:-1] + \ 19 | self[:,0,:,:,1:]) 20 | c_vel_y = torch.zeros_like(self[:,0]) 21 | c_vel_y[:,:,:-1,:] = 0.5 * (self[:,1,:,0:-1,:] + \ 22 | self[:,1,:,1:,:]) 23 | c_vel_z = torch.zeros_like(self[:,0]) 24 | 25 | if (is3D): 26 | c_vel_z = torch.zeros_like(self[:,0]) 27 | c_vel_z[:,:-1,:,:] = 0.5 * (self[:,2,0:-1,:,:] + \ 28 | self[:,2,1:,:,:]) 29 | 30 | return torch.stack((c_vel_x, c_vel_y, c_vel_z), dim=1) 31 | 32 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/init_conditions.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import math 3 | 4 | def createPlumeBCs(batch_dict, density_val, u_scale, rad): 5 | r"""Creates masks to enforce an inlet at the domain bottom wall. 6 | Modifies batch_dict inplace. 7 | Arguments: 8 | batch_dict (dict): Input tensors (p, UDiv, flags, density) 9 | density_val (float): Inlet density. 10 | u_scale (float); Inlet velocity. 11 | rad (float): radius of inlet circle (centered around midpoint of wall) 12 | """ 13 | 14 | cuda = torch.device('cuda') 15 | # batch_dict at input: {p, UDiv, flags, density} 16 | assert len(batch_dict) == 4, "Batch must contain 4 tensors (p, UDiv, flags, density)" 17 | UDiv = batch_dict['U'] 18 | density = batch_dict['density'] 19 | UBC = UDiv.clone().fill_(0) 20 | UBCInvMask = UDiv.clone().fill_(1) 21 | 22 | # Single density value 23 | densityBC = density.clone().fill_(0) 24 | densityBCInvMask = density.clone().fill_(1) 25 | 26 | assert UBC.dim() == 5, 'UBC must have 5 dimensions' 27 | assert UBC.size(0) == 1, 'Only single batches allowed (inference)' 28 | 29 | xdim = UBC.size(4) 30 | ydim = UBC.size(3) 31 | zdim = UBC.size(2) 32 | is3D = (UBC.size(1) == 3) 33 | if not is3D: 34 | assert zdim == 1, 'For 2D, zdim must be 1' 35 | centerX = xdim // 2 36 | centerZ = max( zdim // 2, 1.0) 37 | plumeRad = math.floor(xdim*rad) 38 | 39 | y = 1 40 | if (not is3D): 41 | vec = torch.arange(0,2, device=cuda) 42 | else: 43 | vec = torch.arange(0,3, device=cuda) 44 | vec[2] = 0 45 | 46 | vec.mul_(u_scale) 47 | 48 | index_x = torch.arange(0, xdim, device=cuda).view(xdim).expand_as(density[0][0]) 49 | index_y = torch.arange(0, ydim, device=cuda).view(ydim, 1).expand_as(density[0][0]) 50 | if (is3D): 51 | index_z = torch.arange(0, zdim, device=cuda).view(zdim, 1, 1).expand_as(density[0][0]) 52 | 53 | if (not is3D): 54 | index_ten = torch.stack((index_x, index_y), dim=0) 55 | else: 56 | index_ten = torch.stack((index_x, index_y, index_z), dim=0) 57 | 58 | #TODO 3d implementation 59 | indx_circle = index_ten[:,:,0:4] 60 | indx_circle[0] -= centerX 61 | maskInside = (indx_circle[0].pow(2) <= plumeRad*plumeRad) 62 | 63 | # Inside the plume. Set the BCs. 64 | 65 | #It is clearer to just multiply by mask (casted into Float) 66 | maskInside_f = maskInside.float().clone() 67 | UBC[:,:,:,0:4] = maskInside_f * vec.view(1,2,1,1,1).expand_as(UBC[:,:,:,0:4]).float() 68 | UBCInvMask[:,:,:,0:4].masked_fill_(maskInside, 0) 69 | 70 | densityBC[:,:,:,0:4].masked_fill_(maskInside, density_val) 71 | densityBCInvMask[:,:,:,0:4].masked_fill_(maskInside, 0) 72 | 73 | # Outside the plume. Set the velocity to zero and leave density alone. 74 | 75 | maskOutside = (maskInside == 0) 76 | UBC[:,:,:,0:4].masked_fill_(maskOutside, 0) 77 | UBCInvMask[:,:,:,0:4].masked_fill_(maskOutside, 0) 78 | 79 | # Insert the new tensors in the batch_dict. 80 | batch_dict['UBC'] = UBC 81 | batch_dict['UBCInvMask'] = UBCInvMask 82 | batch_dict['densityBC'] = densityBC 83 | batch_dict['densityBCInvMask'] = densityBCInvMask 84 | 85 | # batch_dict at output = {p, UDiv, flags, density, UBC, 86 | # UBCInvMask, densityBC, densityBCInvMask} 87 | 88 | def createRayleighTaylorBCs(batch_dict, mconf, rho1, rho2): 89 | r"""Creates masks to enforce a Rayleigh-Taylor instability initial conditions. 90 | Top fluid has a density rho1 and lower one rho2. rho1 > rho2 to trigger instability. 91 | Modifies batch_dict inplace. 92 | Arguments: 93 | batch_dict (dict): Input tensors (p, UDiv, flags, density) 94 | mconf (dict): configuration dict (to set thickness and amplitude of interface). 95 | rho1 (float): Top fluid density. 96 | rho2 (float): Lower fluid density. 97 | """ 98 | 99 | cuda = torch.device('cuda') 100 | # batch_dict at input: {p, UDiv, flags, density} 101 | assert len(batch_dict) == 4, "Batch must contain 4 tensors (p, UDiv, flags, density)" 102 | UDiv = batch_dict['U'] 103 | flags = batch_dict['flags'] 104 | 105 | resX = UDiv.size(4) 106 | resY = UDiv.size(3) 107 | 108 | # Here, we just impose initial conditions. 109 | # Upper layer rho2, vel = 0 110 | # Lower layer rho1, vel = 0 111 | 112 | X = torch.arange(0, resX, device=cuda).view(resX).expand((1,resY,resX)) 113 | Y = torch.arange(0, resY, device=cuda).view(resY, 1).expand((1,resY,resX)) 114 | coord = torch.cat((X,Y), dim=0).unsqueeze(0).unsqueeze(2) 115 | 116 | # Atwood number 117 | #A = ((1+rho2) - (1+rho1)) / ((1+rho2) + (1+rho1)) 118 | #print('Atwood number : ' + str(A)) 119 | #density = ((1-A) * torch.tanh(100*(coord[:,1]/resY - (0.85 - \ 120 | # 0.05*torch.cos(math.pi*(coord[:,0]/resX)))))).unsqueeze(1) 121 | thick = mconf['perturbThickness'] 122 | ampl = mconf['perturbAmplitude'] 123 | h = mconf['height'] 124 | density = 0.5*(rho2+rho1 + (rho2-rho1)*torch.tanh(thick*(coord[:,1]/resY - \ 125 | (h + ampl*torch.cos(2*math.pi*(coord[:,0]/resX)))))).unsqueeze(1) 126 | 127 | batch_dict['density'] = density 128 | batch_dict['flags'] = flags 129 | 130 | # batch_dict at output = {p, UDiv, flags, density} 131 | 132 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/set_wall_bcs.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from . import CellType 3 | 4 | def setWallBcs(U, flags): 5 | r"""Sets velocity slip BCs in walls. 6 | 7 | Arguments: 8 | U (Tensor): Input velocity. 9 | flags (Tensor): Input occupancy grid. 10 | Output: 11 | U (Tensor): Output velocity (with enforced BCs). 12 | """ 13 | cuda = torch.device('cuda') 14 | assert (U.dim() == 5 and flags.dim() == 5), 'Dimension mismatch' 15 | assert flags.size(1) == 1, 'flags is not a scalar' 16 | bsz = flags.size(0) 17 | d = flags.size(2) 18 | h = flags.size(3) 19 | w = flags.size(4) 20 | 21 | is3D = (U.size(1) == 3) 22 | if (not is3D): 23 | assert d == 1, '2D velocity field but zdepth > 1' 24 | assert U.size(1) == 2, '2D velocity field must have only 2 channels' 25 | 26 | assert (U.size(0) == bsz and U.size(2) == d and U.size(3) == h and U.size(4) == w),\ 27 | 'Size mismatch' 28 | assert (U.is_contiguous() and flags.is_contiguous()), 'Input is not contiguous' 29 | 30 | i = torch.arange(start=0, end=w, dtype=torch.long, device=cuda) \ 31 | .view(1,w).expand(bsz, d, h, w) 32 | j = torch.arange(start=0, end=h, dtype=torch.long, device=cuda) \ 33 | .view(1,h,1).expand(bsz, d, h, w) 34 | k = torch.zeros_like(i) 35 | if (is3D): 36 | k = torch.arange(start=0, end=d, dtype=torch.long, device=cuda) \ 37 | .view(1,d,1,1).expand(bsz, d, h, w) 38 | 39 | zero = torch.zeros_like(i) 40 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 41 | 42 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 43 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 44 | 45 | mCont = torch.ones_like(zeroBy) 46 | 47 | cur_fluid = flags.eq(CellType.TypeFluid).squeeze(1) 48 | cur_obs = flags.eq(CellType.TypeObstacle).squeeze(1) 49 | mNotFluidNotObs = cur_fluid.ne(1).__and__(cur_obs.ne(1)) 50 | mCont.masked_fill_(mNotFluidNotObs, 0) 51 | 52 | # The neighbour to the left (i-1,j,k) is an obstacle 53 | # Set u.n = u_solid.n (slip bc) (direction i) 54 | i_l = zero.where( (i <=0), i - 1) 55 | #obst100 = zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i_l].eq(CellType.TypeObstacle))).__and__(mCont) 56 | obst100 = flags[idx_b, zero, k, j, i_l].eq(CellType.TypeObstacle).__and__(mCont) 57 | U[:,0].masked_fill_(obst100, 0) 58 | 59 | # Current cell is an obstacle. 60 | # The neighbour to the left (i-1,j,k) is fluid. 61 | # Set u.n = u_solid.n (slip bc) (direction i) 62 | #obs_fluid100 = zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i_l].eq(CellType.TypeFluid))). \ 63 | obs_fluid100 = (flags[idx_b, zero, k, j, i_l].eq(CellType.TypeFluid)). \ 64 | __and__(cur_obs).__and__(mCont) 65 | U[:,0].masked_fill_(obs_fluid100, 0) 66 | 67 | j_l = zero.where( (j <= 0), j - 1) 68 | #obst010 = zeroBy.where( j <= 0, (flags[idx_b, zero, k, j_l, i].eq(CellType.TypeObstacle))).__and__(mCont) 69 | obst010 = (flags[idx_b, zero, k, j_l, i].eq(CellType.TypeObstacle)).__and__(mCont) 70 | U[:,1].masked_fill_(obst010, 0) 71 | #obs_fluid010 = zeroBy.where( j <= 0, (flags[idx_b, zero, k, j_l, i].eq(CellType.TypeFluid))).\ 72 | obs_fluid010 = (flags[idx_b, zero, k, j_l, i].eq(CellType.TypeFluid)).\ 73 | __and__(cur_obs).__and__(mCont) 74 | U[:,1].masked_fill_(obs_fluid010, 0) 75 | 76 | if (is3D): 77 | k_l = zero.where( (k <= 0), k - 1) 78 | 79 | obst001 = zeroBy.where( k <= 0, (flags[idx_b, zero, k_l, j, i].eq(CellType.TypeObstacle))).__and__(mCont) 80 | U[:,2].masked_fill_(obst001, 0) 81 | 82 | obs_fluid001 = zeroBy.where( k <= 0, (flags[idx_b, zero, k_l, j, i].eq(CellType.TypeFluid))). \ 83 | _and__(cur_obs).__and__(mCont) 84 | U[:,2].masked_fill_(obs_fluid001, 0) 85 | 86 | return U 87 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/set_wall_bcs_inflow.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def setWallBcs(U, flags): 4 | 5 | cuda = torch.device('cuda') 6 | assert (U.dim() == 5 and flags.dim() == 5), 'Dimension mismatch' 7 | assert flags.size(1) == 1, 'flags is not a scalar' 8 | bsz = flags.size(0) 9 | d = flags.size(2) 10 | h = flags.size(3) 11 | w = flags.size(4) 12 | 13 | is3D = (U.size(1) == 3) 14 | if (not is3D): 15 | assert d == 1, '2D velocity field but zdepth > 1' 16 | assert U.size(1) == 2, '2D velocity field must have only 2 channels' 17 | 18 | assert (U.size(0) == bsz and U.size(2) == d and U.size(3) == h and U.size(4) == w),\ 19 | 'Size mismatch' 20 | assert (U.is_contiguous() and flags.is_contiguous()), 'Input is not contiguous' 21 | 22 | TypeFluid = 1 23 | TypeObstacle = 2 24 | TypeInflow = 8 25 | 26 | i = torch.arange(start=0, end=w, dtype=torch.long, device=cuda) \ 27 | .view(1,w).expand(bsz, d, h, w) 28 | j = torch.arange(start=0, end=h, dtype=torch.long, device=cuda) \ 29 | .view(1,h,1).expand(bsz, d, h, w) 30 | k = torch.zeros_like(i) 31 | if (is3D): 32 | k = torch.arange(start=0, end=d, dtype=torch.long, device=cuda) \ 33 | .view(1,d,1,1).expand(bsz, d, h, w) 34 | 35 | zero = torch.zeros_like(i) 36 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 37 | 38 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 39 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 40 | 41 | mCont = torch.ones_like(zeroBy) 42 | 43 | cur_fluid = flags.eq(TypeFluid).squeeze(1) 44 | cur_obs = flags.eq(TypeObstacle).squeeze(1) 45 | mNotFluidNotObs = cur_fluid.ne(1).__and__(cur_obs.ne(1)) 46 | mCont.masked_fill_(mNotFluidNotObs, 0) 47 | 48 | # The neighbour to the left (i-1,j,k) is an obstacle 49 | # Set u.n = u_solid.n (slip bc) (direction i) 50 | i_l = zero.where( (i <=0), i - 1) 51 | obst100 = zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i_l].eq(TypeObstacle))).__and__(mCont) 52 | U[:,0].masked_fill_(obst100, 0) 53 | 54 | # Current cell is an obstacle. 55 | # The neighbour to the left (i-1,j,k) is fluid. 56 | # Set u.n = u_solid.n (slip bc) (direction i) 57 | obs_fluid100 = zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i_l].eq(TypeFluid))). \ 58 | __and__(cur_obs).__and__(mCont) 59 | U[:,0].masked_fill_(obs_fluid100, 0) 60 | 61 | j_l = zero.where( (j <= 0), j - 1) 62 | obst010 = zeroBy.where( j <= 0, (flags[idx_b, zero, k, j_l, i].eq(TypeObstacle))).__and__(mCont) 63 | U[:,1].masked_fill_(obst010, 0) 64 | obs_fluid010 = zeroBy.where( j <= 0, (flags[idx_b, zero, k, j_l, i].eq(TypeFluid))).\ 65 | __and__(cur_obs).__and__(mCont) 66 | U[:,1].masked_fill_(obs_fluid010, 0) 67 | 68 | if (is3D): 69 | k_l = zero.where( (k <= 0), k - 1) 70 | 71 | obst001 = zeroBy.where( k <= 0, (flags[idx_b, zero, k_l, j, i].eq(TypeObstacle))).__and__(mCont) 72 | U[:,2].masked_fill_(obst001, 0) 73 | 74 | obs_fluid001 = zeroBy.where( k <= 0, (flags[idx_b, zero, k_l, j, i].eq(TypeFluid))). \ 75 | _and__(cur_obs).__and__(mCont) 76 | U[:,2].masked_fill_(obs_fluid001, 0) 77 | 78 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/set_wall_bcs_stick.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import CellType 4 | 5 | def setWallBcsStick(U, flags, flags_stick): 6 | 7 | cuda = torch.device('cuda') 8 | assert (U.dim() == 5 and flags.dim() == 5 and flags_stick.dim() == 5), 'Dimension mismatch' 9 | assert flags.size(1) == 1, 'flags is not a scalar' 10 | assert flags_stick.size(1) == 1, 'flags is not a scalar' 11 | bsz = flags.size(0) 12 | d = flags.size(2) 13 | h = flags.size(3) 14 | w = flags.size(4) 15 | 16 | is3D = (U.size(1) == 3) 17 | if (not is3D): 18 | assert d == 1, '2D velocity field but zdepth > 1' 19 | assert U.size(1) == 2, '2D velocity field must have only 2 channels' 20 | 21 | assert (U.size(0) == bsz and U.size(2) == d and U.size(3) == h and U.size(4) == w),\ 22 | 'Size mismatch' 23 | assert (U.is_contiguous() and flags.is_contiguous() and flags_stick.is_contiguous()), 'Input is not contiguous' 24 | 25 | i = torch.arange(start=0, end=w, dtype=torch.long, device=cuda) \ 26 | .view(1,w).expand(bsz, d, h, w) 27 | j = torch.arange(start=0, end=h, dtype=torch.long, device=cuda) \ 28 | .view(1,h,1).expand(bsz, d, h, w) 29 | k = torch.zeros_like(i) 30 | if (is3D): 31 | k = torch.arange(start=0, end=d, dtype=torch.long, device=cuda) \ 32 | .view(1,d,1,1).expand(bsz, d, h, w) 33 | 34 | zero = torch.zeros_like(i) 35 | maxX = torch.zeros_like(i).fill_(w-1) 36 | maxY = torch.zeros_like(i).fill_(h-1) 37 | zeroF = zero.float() 38 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 39 | 40 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 41 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 42 | 43 | mCont = torch.ones_like(zeroBy) 44 | 45 | cur_fluid = flags.eq(CellType.TypeFluid).squeeze(1) 46 | cur_obs = flags.eq(CellType.TypeObstacle).squeeze(1) 47 | cur_stick = flags_stick.eq(CellType.TypeStick).squeeze(1) 48 | 49 | stick_is_obs = cur_obs.masked_select(cur_stick.eq(CellType.TypeStick).squeeze(1)) 50 | assert stick_is_obs.all(), 'Stick cells must be also obstacles.' 51 | mNotCells = cur_fluid.ne(1).__and__\ 52 | (cur_obs.ne(1)).__and__\ 53 | (cur_stick.ne(1)) 54 | mCont.masked_fill_(mNotCells, 0) 55 | 56 | # First, set all velocities INSIDE obstacles to zero. 57 | U.masked_fill_(cur_obs.unsqueeze(1), 0) 58 | 59 | # Slip Boundary Conditon. Normal vel is 0. 60 | # The neighbour to the left (i-1,j,k) is an obstacle 61 | im1 = zero.where(i <=0, i - 1) 62 | obst_im1jk = zeroBy.where(i <= 0, (flags[idx_b, zero, k, j, im1].eq(TypeObstacle))).__and__(mCont) 63 | U[:,0].masked_fill_(obst_im1jk, 0) 64 | 65 | # Current cell is an obstacle. 66 | # The neighbour to the left (i-1,j,k) is fluid. 67 | # Set normal direction velocity to 0. 68 | obs_ijk_fluid_im1jk = zeroBy.where(i <= 0, (flags[idx_b, zero, k, j, im1].eq(TypeFluid))). \ 69 | __and__(cur_obs).__and__(mCont) 70 | U[:,0].masked_fill_(obs_ijk_fluid_im1jk, 0) 71 | 72 | jm1 = zero.where(j <= 0, j - 1) 73 | obst_ijm1k = zeroBy.where(j <= 0, (flags[idx_b, zero, k, jm1, i].eq(TypeObstacle))).__and__(mCont) 74 | U[:,1].masked_fill_(obst_ijm1k, 0) 75 | obs_ijk_fluid_ijm1k = zeroBy.where(j <= 0, (flags[idx_b, zero, k, jm1, i].eq(TypeFluid))).\ 76 | __and__(cur_obs).__and__(mCont) 77 | U[:,1].masked_fill_(obs_ijk_fluid_ijm1k, 0) 78 | 79 | if (is3D): 80 | km1 = zero.where(k <= 0, k - 1) 81 | 82 | obst_ijkm1 = zeroBy.where(k <= 0, (flags[idx_b, zero, km1, j, i].eq(TypeObstacle))).__and__(mCont) 83 | U[:,2].masked_fill_(obst_ijkm1, 0) 84 | 85 | obs_ijk_fluid_ijkm1 = zeroBy.where(k <= 0, (flags[idx_b, zero, km1, j, i].eq(TypeFluid))). \ 86 | _and__(cur_obs).__and__(mCont) 87 | U[:,2].masked_fill_(obs_ijk_fluid_ijkm1_, 0) 88 | 89 | # No-slip (aka stick) Boundary condition. 90 | # Normal AND tangential velocities are zero. 91 | # As the stick cells are also obstacle, we just need to add 92 | # tangential vel=0. 93 | # We work only on ghost cells (a velocity in obstacle cell) to enforce this condition. 94 | # For that reason, and needing access to multiple velocities, let's operate with float masks. 95 | 96 | ip1 = maxX.where(i>=(w-1), i + 1) 97 | jp1 = maxY.where(j>=(h-1), j + 1) 98 | 99 | cur_stick = cur_stick 100 | 101 | # Vertical velocities: 102 | fluid_im1jk = (flags[idx_b, zero, k, j, im1].eq(TypeFluid)).__and__(mCont) 103 | fluid_ip1jk = (flags[idx_b, zero, k, j, ip1].eq(TypeFluid)).__and__(mCont) 104 | 105 | v_im1jk = zeroF.where(i<=0, (U[idx_b, 1, k, j, im1])) 106 | v_ip1jk = zeroF.where(i>=(w-1), (U[idx_b, 1, k, j, ip1])) 107 | # Current cell is stick and left neighbor fluid. v(i,j) = -V(i-1,j) 108 | ghost_ijk_fluid_im1jk = cur_stick.__and__(fluid_im1jk).__and__(mCont) 109 | U[:,1].masked_scatter_(ghost_ijk_fluid_im1jk, \ 110 | (-v_im1jk).masked_select(ghost_ijk_fluid_im1jk)) 111 | # Current cell is stick and right neighbor fluid. v(i,j) = -V(i+1,j) 112 | ghost_ijk_fluid_ip1jk = cur_stick.__and__(fluid_ip1jk).__and__(mCont) 113 | U[:,1].masked_scatter_(ghost_ijk_fluid_ip1jk, \ 114 | (-v_ip1jk).masked_select(ghost_ijk_fluid_ip1jk)) 115 | # Both neighbors left and right are fluid. We approximate the ghost velocity as the mean between the two velocities. 116 | # This case should be avoided as much as possible (by putting walls of thickness 2 at least). 117 | ghost_ijk_fluid_imp1jk = cur_stick.__and__(fluid_im1jk).__and__(fluid_ip1jk).__and__(mCont) 118 | U[:,1].masked_scatter_(ghost_ijk_fluid_imp1jk, \ 119 | ((0.5)*(-v_im1jk-v_ip1jk)).masked_select(ghost_ijk_fluid_imp1jk)) 120 | 121 | # For horizontal velocities: 122 | fluid_ijm1k = (flags[idx_b, zero, k, jm1, i].eq(TypeFluid)).__and__(mCont) 123 | fluid_ijp1k = (flags[idx_b, zero, k, jp1, i].eq(TypeFluid)).__and__(mCont) 124 | 125 | u_ijm1k = zeroF.where(j<=0, (U[idx_b, 0, k, jm1, i])) 126 | u_ijp1k = zeroF.where(j>=(h-1), (U[idx_b, 0, k, jp1, i])) 127 | 128 | # Current cell is stick and bottom neighbor fluid. u(i,j) = -u(i,j-1) 129 | ghost_ijk_fluid_ijm1k = cur_stick.__and__(fluid_ijm1k).__and__(mCont) 130 | U[:,0].masked_scatter_(ghost_ijk_fluid_ijm1k, \ 131 | (-u_ijm1k).masked_select(ghost_ijk_fluid_ijm1k)) 132 | # Current cell is stick and upper neighbor fluid. u(i,j) = -u(i,j+1) 133 | ghost_ijk_fluid_ijp1k = cur_stick.__and__(fluid_ijp1k).__and__(mCont) 134 | U[:,0].masked_scatter_(ghost_ijk_fluid_ijp1k, \ 135 | (-u_ijp1k).masked_select(ghost_ijk_fluid_ijp1k)) 136 | # Both neighbors up and bottom are fluid. We approximate the ghost velocity as the mean between the two velocities. 137 | # This case should be avoided as much as possible (by putting walls of thickness 2 at least). 138 | ghost_ijk_fluid_ijmp1k = cur_stick.__and__(fluid_ijm1k).__and__(fluid_ijm1k).__and__(mCont) 139 | U[:,0].masked_scatter_(ghost_ijk_fluid_ijmp1k, \ 140 | ((0.5)*(-u_ijm1k-u_ijp1k)).masked_select(ghost_ijk_fluid_ijmp1k)) 141 | 142 | # Corner cases: 143 | cur_stick = cur_stick.float() 144 | left_stick = (flags_stick[idx_b, zero, k, j, im1].eq(TypeStick)).__and__(mCont) 145 | right_stick = (flags_stick[idx_b, zero, k, j, ip1].eq(TypeStick)).__and__(mCont) 146 | bottom_stick = (flags_stick[idx_b, zero, k, jm1, i].eq(TypeStick)).__and__(mCont) 147 | upper_stick = (flags_stick[idx_b, zero, k, jp1, i].eq(TypeStick)).__and__(mCont) 148 | 149 | # U velocity. 150 | G_U_ijk = (cur_stick.float() + left_stick.float() + bottom_stick.float() + \ 151 | cur_stick.float() + left_stick.float() + upper_stick.float()).eq(3) 152 | U[:,0].masked_fill_(G_U_ijk, 0) 153 | 154 | # V velocity. 155 | G_V_ijk = (cur_stick.float() + left_stick.float() + bottom_stick.float() + \ 156 | cur_stick.float() + right_stick.float() + bottom_stick.float()).eq(3) 157 | U[:,1].masked_fill_(G_V_ijk, 0) 158 | 159 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/source_terms.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import CellType 4 | from . import getDx 5 | 6 | def addBuoyancy(U, flags, density, gravity, rho_star, dt): 7 | r"""Add buoyancy force. 8 | Arguments: 9 | U (Tensor): velocity field (size(2) can be 2 or 3, indicating 2D / 3D) 10 | flags (Tensor): input occupancy grid. 11 | density (Tensor): scalar density grid. 12 | gravity (Tensor): 3D vector indicating direction of gravity. 13 | dt (float): scalar timestep. 14 | Output: 15 | U (Tensor): Output velocity 16 | """ 17 | cuda = torch.device('cuda') 18 | # Argument check 19 | assert U.dim() == 5 and flags.dim() == 5 and density.dim() == 5,\ 20 | "Dimension mismatch" 21 | assert flags.size(1) == 1, "flags is not scalar" 22 | bsz = flags.size(0) 23 | d = flags.size(2) 24 | h = flags.size(3) 25 | w = flags.size(4) 26 | 27 | is3D = (U.size(1) == 3) 28 | 29 | bnd = 1 30 | 31 | if not is3D: 32 | assert d == 1, "2D velocity field but zdepth > 1" 33 | assert U.size(1) == 2, "2D velocity field must have only 2 channels" 34 | 35 | assert U.size(0) == bsz and U.size(2) == d and \ 36 | U.size(3) == h and U.size(4) == w, "Size mismatch" 37 | assert density.is_same_size(flags), "Size mismatch" 38 | 39 | assert U.is_contiguous() and flags.is_contiguous() and \ 40 | density.is_contiguous(), "Input is not contiguous" 41 | 42 | assert gravity.dim() == 1 and gravity.size(0) == 3, \ 43 | "Gravity must be a 3D vector (even in 2D)" 44 | 45 | # (aalgua) I don't know why Manta divides by dx, as in all other modules 46 | # dx = 1. 47 | strength = gravity * dt 48 | 49 | i = torch.arange(0, w, dtype=torch.long, device=cuda).view(1,w).expand(bsz, d, h, w) 50 | j = torch.arange(0, h, dtype=torch.long, device=cuda).view(1,h,1).expand(bsz, d, h, w) 51 | k = torch.zeros_like(i) 52 | if (is3D): 53 | k = torch.arange(0, d, dtype=torch.long, device=cuda).view(1,d,1,1).expand(bsz, d, h, w) 54 | 55 | zero = torch.zeros_like(i) 56 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 57 | zero_f = zero.cuda().float() 58 | 59 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 60 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 61 | 62 | maskBorder = (i < bnd).__or__\ 63 | (i > w - 1 - bnd).__or__\ 64 | (j < bnd).__or__\ 65 | (j > h - 1 - bnd) 66 | if (is3D): 67 | maskBorder = maskBorder.__or__(k < bnd).__or__\ 68 | (k > d - 1 - bnd) 69 | 70 | maskBorder = maskBorder.unsqueeze(1) 71 | 72 | # No buoyancy on the border. Set continue (mCont) to false. 73 | mCont = torch.ones_like(zeroBy).unsqueeze(1) 74 | mCont.masked_fill_(maskBorder, 0) 75 | 76 | isFluid = flags.eq(CellType.TypeFluid).__and__(mCont) 77 | mCont.masked_fill_(isFluid.ne(1), 0) 78 | mCont.squeeze_(1) 79 | 80 | max_X = torch.zeros_like(zero).fill_(w-1) 81 | max_Y = torch.zeros_like(zero).fill_(h-1) 82 | 83 | i_l = zero.where( (i <= 0), i-1) 84 | i_r = max_X.where( (i > w-2), i+1) 85 | j_l = zero.where( (j <= 0), j-1) 86 | j_r = max_Y.where( (j > h-2), j+1) 87 | 88 | rho_star = rho_star 89 | 90 | fluid100 = flags[idx_b, zero, k, j, i_l].eq(CellType.TypeFluid).__and__ \ 91 | (mCont) 92 | #(flags[idx_b, zero, k, j, i_r].eq(CellType.TypeFluid)).__and__ \ 93 | 94 | factor = strength[0] * ((0.5* \ 95 | (density[idx_b, zero, k, j, i] + 96 | density[idx_b, zero, k, j, i_l] ) - rho_star)) 97 | U[:,0].masked_scatter_(fluid100, (U.select(1,0) + factor).masked_select(fluid100)) 98 | 99 | fluid010 = flags[idx_b, zero, k, j_l, i].eq(CellType.TypeFluid).__and__ \ 100 | (mCont) 101 | #(flags[idx_b, zero, k, j_r, i].eq(CellType.TypeFluid)).__and__ \ 102 | #fluid010 = zeroBy.where( j <= 0, (flags[idx_b, zero, k, j-1, i].eq(CellType.TypeFluid))).__and__(mCont) 103 | factor = strength[1] * ((0.5* \ 104 | (density[idx_b, zero, k, j, i] + 105 | density[idx_b, zero, k, j_l, i] ) - rho_star)) 106 | #factor = strength[1] * (density.squeeze(1) - \ 107 | # zero_f.where( j <= 0, density[idx_b, zero, k, j-1, i]) ) 108 | U[:,1].masked_scatter_(fluid010, (U.select(1,1) + factor).masked_select(fluid010)) 109 | 110 | if (is3D): 111 | fluid001 = zeroBy.where( j <= 0, (flags[idx_b, zero, k-1, j, i].eq(CellType.TypeFluid))).__and__(mCont) 112 | factor = strength[2] *(0.5* (density.squeeze(1) + \ 113 | zero_f.where(k <= 1, density[idx_b, zero, k-1, j, i]) )) 114 | U[:,2].masked_scatter_(fluid001, (U.select(1,2) + factor).masked_select(fluid001)) 115 | 116 | return U 117 | 118 | # ***************************************************************************** 119 | # addGravity 120 | # ***************************************************************************** 121 | 122 | def addGravity(U, flags, gravity, dt): 123 | r"""Add gravity force. 124 | 125 | Arguments: 126 | U (Tensor): velocity field (size(2) can be 2 or 3, indicating 2D / 3D) 127 | flags (Tensor): input occupancy grid. 128 | gravity (Tensor): 3D vector indicating direction of gravity. 129 | dt (float): scalar timestep. 130 | Output: 131 | U (Tensor): Output velocity 132 | """ 133 | cuda = torch.device('cuda') 134 | # Argument check 135 | assert U.dim() == 5 and flags.dim() == 5, "Dimension mismatch" 136 | assert flags.size(1) == 1, "flags is not scalar" 137 | bsz = flags.size(0) 138 | d = flags.size(2) 139 | h = flags.size(3) 140 | w = flags.size(4) 141 | 142 | is3D = (U.size(1) == 3) 143 | 144 | bnd = 1 145 | if not is3D: 146 | assert d == 1, "2D velocity field but zdepth > 1" 147 | assert U.size(1) == 2, "2D velocity field must have only 2 channels" 148 | 149 | assert U.size(0) == bsz and U.size(2) == d and \ 150 | U.size(3) == h and U.size(4) == w, "Size mismatch" 151 | 152 | assert U.is_contiguous() and flags.is_contiguous(), "Input is not contiguous" 153 | 154 | assert gravity.dim() == 1 and gravity.size(0) == 3,\ 155 | "Gravity must be a 3D vector (even in 2D)" 156 | 157 | # (aalgua) I don't know why Manta divides by dx, as in all other modules 158 | # dx = 1. 159 | force = gravity * dt 160 | 161 | i = torch.arange(0, w, dtype=torch.long, device=cuda).view(1,w).expand(bsz, d, h, w) 162 | j = torch.arange(0, h, dtype=torch.long, device=cuda).view(1,h,1).expand(bsz, d, h, w) 163 | k = torch.zeros_like(i) 164 | if (is3D): 165 | k = torch.arange(0, d, dtype=torch.long, device=cuda).view(1,d,1,1).expand(bsz, d, h, w) 166 | 167 | zero = torch.zeros_like(i) 168 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 169 | zero_f = zero.float() 170 | 171 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 172 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 173 | 174 | maskBorder = (i < bnd).__or__\ 175 | (i > w - 1 - bnd).__or__\ 176 | (j < bnd).__or__\ 177 | (j > h - 1 - bnd) 178 | if (is3D): 179 | maskBorder = maskBorder.__or__(k < bnd).__or__(k > d - 1 - bnd) 180 | 181 | maskBorder = maskBorder.unsqueeze(1) 182 | 183 | # No buoyancy on the border. Set continue (mCont) to false. 184 | mCont = torch.ones_like(zeroBy).unsqueeze(1) 185 | mCont.masked_fill_(maskBorder, 0) 186 | 187 | cur_fluid = flags.eq(CellType.TypeFluid).__and__(mCont) 188 | cur_empty = flags.eq(CellType.TypeEmpty).__and__(mCont) 189 | 190 | mNotFluidNotEmpt = cur_fluid.ne(1).__and__(cur_empty.ne(1)) 191 | mCont.masked_fill_(mNotFluidNotEmpt, 0) 192 | 193 | mCont.squeeze_(1) 194 | #print() 195 | #print('F = ') 196 | #print(force) 197 | 198 | #print('before') 199 | #print(U) 200 | #print(U.size()) 201 | fluid100 = (zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i-1].eq(CellType.TypeFluid))) \ 202 | .__or__(( zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i-1].eq(CellType.TypeEmpty)))) \ 203 | .__and__(cur_fluid.squeeze(1)))).__and__(mCont) 204 | U[:,0].masked_scatter_(fluid100, (U[:,0] + force[0]).masked_select(fluid100)) 205 | 206 | fluid010 = (zeroBy.where( j <= 0, (flags[idx_b, zero, k, j-1, i].eq(CellType.TypeFluid))) \ 207 | .__or__(( zeroBy.where( j <= 0, (flags[idx_b, zero, k, j-1, i].eq(CellType.TypeEmpty)))) \ 208 | .__and__(cur_fluid.squeeze(1))) ).__and__(mCont) 209 | U[:,1].masked_scatter_(fluid010, (U[:,1] + force[1]).masked_select(fluid010)) 210 | 211 | if (is3D): 212 | fluid001 = (zeroBy.where( k <= 0, (flags[idx_b, zero, k-1, j, i].eq(CellType.TypeFluid))) \ 213 | .__or__(( zeroBy.where( k <= 0, (flags[idx_b, zero, k-1, j, i].eq(CellType.TypeEmpty)))) \ 214 | .__and__(cur_fluid.squeeze(1)))).__and__(mCont) 215 | U[:,2].masked_scatter_(fluid001, (U[:,2] + force[2]).masked_select(fluid001)) 216 | #print('after') 217 | #print(U) 218 | #print(U.size()) 219 | return U 220 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/source_terms_test.py: -------------------------------------------------------------------------------- 1 | from . import getDx 2 | import torch 3 | 4 | # ***************************************************************************** 5 | # addBuoyancy 6 | # ***************************************************************************** 7 | 8 | # Add buoyancy force. AddBuoyancy has a dt term. 9 | # Note: Buoyancy is added IN-PLACE. 10 | # 11 | # @input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 12 | # @input flags - input occupancy grid 13 | # @input density - scalar density grid. 14 | # @input gravity - 3D vector indicating direction of gravity. 15 | # @input dt - scalar timestep. 16 | 17 | def addBuoyancy(U, flags, density, gravity, dt): 18 | cuda = torch.device('cuda') 19 | # Argument check 20 | assert U.dim() == 5 and flags.dim() == 5 and density.dim() == 5,\ 21 | "Dimension mismatch" 22 | assert flags.size(1) == 1, "flags is not scalar" 23 | bsz = flags.size(0) 24 | d = flags.size(2) 25 | h = flags.size(3) 26 | w = flags.size(4) 27 | 28 | is3D = (U.size(1) == 3) 29 | 30 | bnd = 1 31 | 32 | if not is3D: 33 | assert d == 1, "2D velocity field but zdepth > 1" 34 | assert U.size(1) == 2, "2D velocity field must have only 2 channels" 35 | 36 | assert U.size(0) == bsz and U.size(2) == d and \ 37 | U.size(3) == h and U.size(4) == w, "Size mismatch" 38 | assert density.is_same_size(flags), "Size mismatch" 39 | 40 | assert U.is_contiguous() and flags.is_contiguous() and \ 41 | density.is_contiguous(), "Input is not contiguous" 42 | 43 | assert gravity.dim() == 1 and gravity.size(0) == 3, \ 44 | "Gravity must be a 3D vector (even in 2D)" 45 | 46 | TypeFluid = 1 47 | 48 | # (aalgua) I don't know why Manta divides by dx, as in all other modules 49 | # dx = 1. I will leave it, but the input gravity should be multiplied 50 | # by getDx(flags) 51 | strength = - gravity * dt 52 | 53 | i = torch.arange(0, w, dtype=torch.long, device=cuda).view(1,w).expand(bsz, d, h, w) 54 | j = torch.arange(0, h, dtype=torch.long, device=cuda).view(1,h,1).expand(bsz, d, h, w) 55 | k = torch.zeros_like(i) 56 | if (is3D): 57 | k = torch.arange(0, d, dtype=torch.long, device=cuda).view(1,d,1,1).expand(bsz, d, h, w) 58 | 59 | zero = torch.zeros_like(i) 60 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 61 | zero_f = zero.cuda().float() 62 | 63 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 64 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 65 | 66 | 67 | maskBorder = (i < bnd).__or__\ 68 | (i > w - 1 - bnd).__or__\ 69 | (j < bnd).__or__\ 70 | (j > h - 1 - bnd) 71 | if (is3D): 72 | maskBorder = maskBorder.__or__(k < bnd).__or__\ 73 | (k > d - 1 - bnd) 74 | 75 | maskBorder = maskBorder.unsqueeze(1) 76 | 77 | # No buoyancy on the border. Set continue (mCont) to false. 78 | mCont = torch.ones_like(zeroBy).unsqueeze(1) 79 | mCont.masked_fill_(maskBorder, 0) 80 | 81 | isFluid = flags.eq(TypeFluid).__and__(mCont) 82 | mCont.masked_fill_(isFluid.ne(1), 0) 83 | mCont.squeeze_(1) 84 | 85 | fluid100 = zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i-1].eq(TypeFluid))).__and__(mCont) 86 | factor = 0.5 * strength[0] * (density.squeeze(1) + \ 87 | zero_f.where(i <= 0, density[idx_b, zero, k, j, i-1]) ) 88 | U[:,0].masked_scatter_(fluid100, (U.select(1,0) + factor).masked_select(fluid100)) 89 | 90 | fluid010 = zeroBy.where( j <= 0, (flags[idx_b, zero, k, j-1, i].eq(TypeFluid))).__and__(mCont) 91 | factor = 0.5 * strength[1] * (density.squeeze(1) + \ 92 | zero_f.where( j <= 0, density[idx_b, zero, k, j-1, i]) ) 93 | U[:,1].masked_scatter_(fluid010, (U.select(1,1) + factor).masked_select(fluid010)) 94 | 95 | if (is3D): 96 | fluid001 = zeroBy.where( j <= 0, (flags[idx_b, zero, k-1, j, i].eq(TypeFluid))).__and__(mCont) 97 | factor = 0.5 * strength[2] * (density.squeeze(1) + \ 98 | zero_f.where(k <= 1, density[idx_b, zero, k-1, j, i]) ) 99 | U[:,2].masked_scatter_(fluid001, (U.select(1,2) + factor).masked_select(fluid001)) 100 | 101 | 102 | 103 | # ***************************************************************************** 104 | # addGravity 105 | # ***************************************************************************** 106 | 107 | # Add gravity force. It has a dt term. 108 | # Note: gravity is added IN-PLACE. 109 | # 110 | # @input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 111 | # @input flags - input occupancy grid 112 | # @input gravity - 3D vector indicating direction of gravity. 113 | # @input dt - scalar timestep. 114 | 115 | def addGravity(U, flags, gravity, dt): 116 | 117 | cuda = torch.device('cuda') 118 | # Argument check 119 | assert U.dim() == 5 and flags.dim() == 5, "Dimension mismatch" 120 | assert flags.size(1) == 1, "flags is not scalar" 121 | bsz = flags.size(0) 122 | d = flags.size(2) 123 | h = flags.size(3) 124 | w = flags.size(4) 125 | 126 | is3D = (U.size(1) == 3) 127 | 128 | bnd = 1 129 | if not is3D: 130 | assert d == 1, "2D velocity field but zdepth > 1" 131 | assert U.size(1) == 2, "2D velocity field must have only 2 channels" 132 | 133 | assert U.size(0) == bsz and U.size(2) == d and \ 134 | U.size(3) == h and U.size(4) == w, "Size mismatch" 135 | 136 | assert U.is_contiguous() and flags.is_contiguous(), "Input is not contiguous" 137 | 138 | assert gravity.dim() == 1 and gravity.size(0) == 3,\ 139 | "Gravity must be a 3D vector (even in 2D)" 140 | 141 | TypeFluid = 1 142 | TypeObstacle = 2 143 | TypeEmpty = 4 144 | 145 | # (aalgua) I don't know why Manta divides by dx, as in all other modules 146 | # dx = 1. I will leave it, but the input gravity should be multiplied 147 | # by getDx(flags) 148 | force = gravity * dt 149 | 150 | i = torch.arange(0, w, dtype=torch.long, device=cuda).view(1,w).expand(bsz, d, h, w) 151 | j = torch.arange(0, h, dtype=torch.long, device=cuda).view(1,h,1).expand(bsz, d, h, w) 152 | k = torch.zeros_like(i) 153 | if (is3D): 154 | k = torch.arange(0, d, dtype=torch.long, device=cuda).view(1,d,1,1).expand(bsz, d, h, w) 155 | 156 | zero = torch.zeros_like(i) 157 | zeroBy = torch.zeros(i.size(), dtype=torch.uint8, device=cuda) 158 | zero_f = zero.float() 159 | 160 | idx_b = torch.arange(start=0, end=bsz, dtype=torch.long, device=cuda) \ 161 | .view(bsz, 1, 1, 1).expand(bsz,d,h,w) 162 | 163 | maskBorder = (i < bnd).__or__\ 164 | (i > w - 1 - bnd).__or__\ 165 | (j < bnd).__or__\ 166 | (j > h - 1 - bnd) 167 | if (is3D): 168 | maskBorder = maskBorder.__or__(k < bnd).__or__(k > d - 1 - bnd) 169 | 170 | maskBorder = maskBorder.unsqueeze(1) 171 | 172 | # No buoyancy on the border. Set continue (mCont) to false. 173 | mCont = torch.ones_like(zeroBy).unsqueeze(1) 174 | mCont.masked_fill_(maskBorder, 0) 175 | 176 | cur_fluid = flags.eq(TypeFluid).__and__(mCont) 177 | cur_empty = flags.eq(TypeEmpty).__and__(mCont) 178 | 179 | mNotFluidNotEmpt = cur_fluid.ne(1).__and__(cur_empty.ne(1)) 180 | mCont.masked_fill_(mNotFluidNotEmpt, 0) 181 | 182 | mCont.squeeze_(1) 183 | 184 | fluid100 = (zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i-1].eq(TypeFluid))) \ 185 | .__or__(( zeroBy.where( i <= 0, (flags[idx_b, zero, k, j, i-1].eq(TypeEmpty)))) \ 186 | .__and__(cur_fluid.squeeze(1)))).__and__(mCont) 187 | U[:,0].masked_scatter_(fluid100, (U.select(1,0) + force[0]).masked_select(fluid100)) 188 | 189 | fluid010 = (zeroBy.where( j <= 0, (flags[idx_b, zero, k, j-1, i].eq(TypeFluid))) \ 190 | .__or__(( zeroBy.where( j <= 0, (flags[idx_b, zero, k, j-1, i].eq(TypeEmpty)))) \ 191 | .__and__(cur_fluid.squeeze(1))) ).__and__(mCont) 192 | U[:,1].masked_scatter_(fluid010, (U.select(1,1) + force[1]).masked_select(fluid010)) 193 | 194 | if (is3D): 195 | fluid001 = (zeroBy.where( k <= 0, (flags[idx_b, zero, k-1, j, i].eq(TypeFluid))) \ 196 | .__or__(( zeroBy.where( k <= 0, (flags[idx_b, zero, k-1, j, i].eq(TypeEmpty)))) \ 197 | .__and__(cur_fluid.squeeze(1)))).__and__(mCont) 198 | U[:,2].masked_scatter_(fluid001, (U.select(1,2) + force[2]).masked_select(fluid001)) 199 | 200 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/util.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import CellType 4 | 5 | def emptyDomain(flags, boundary_width = 1): 6 | r"""Sets all inner cells to Fluid Type flag and domain border to Obstacle Type. 7 | Use this to initialize cleanly the flag Tensor 8 | Arguments: 9 | flags (Tensor): Input flags grid (modified inplace) 10 | boundary_width (int, optional): Set width of border wall. Default set to 1. 11 | """ 12 | 13 | cuda = torch.device('cuda') 14 | assert boundary_width > 0, 'Boundary width must be greater than zero!' 15 | assert flags.dim() == 5, 'Flags tensor should be 5D' 16 | assert flags.size(1) == 1, 'Flags should have only one channels (scalar field)' 17 | 18 | is3D = flags.size(2) > 1 19 | 20 | assert ( ((not is3D) or (flags.size(2) > boundary_width*2)) and \ 21 | (flags.size(3) > boundary_width*2) or (flags.size(4) > boundary_width*2)), \ 22 | 'Simulation domain is not big enough' 23 | xdim = flags.size(4) 24 | ydim = flags.size(3) 25 | zdim = flags.size(2) 26 | bnd = boundary_width 27 | 28 | index_x = torch.arange(0, xdim, device=cuda).view(xdim).expand_as(flags[0][0]) 29 | index_y = torch.arange(0, ydim, device=cuda).view(ydim, 1).expand_as(flags[0][0]) 30 | if (is3D): 31 | index_z = torch.arange(0, zdim, device=cuda).view(zdim, 1 , 1).expand_as(flags[0][0]) 32 | 33 | if (not is3D): 34 | index_ten = torch.stack((index_x, index_y), dim=0) 35 | else: 36 | index_ten = torch.stack((index_x, index_y, index_z), dim=0) 37 | 38 | maskBorder = (index_ten.select(0,0) < bnd).__or__ \ 39 | (index_ten.select(0,0) > xdim - 1 - bnd).__or__\ 40 | (index_ten.select(0,1) < bnd).__or__\ 41 | (index_ten.select(0,1) > ydim - 1 - bnd) 42 | if (is3D): 43 | maskBorder = maskBorder.__or__(index_ten.select(0,2) < bnd).__or__\ 44 | (index_ten.select(0,2) > zdim - 1 - bnd) 45 | 46 | flags.masked_fill_(maskBorder, CellType.TypeObstacle) 47 | flags.masked_fill_((maskBorder == 0), CellType.TypeFluid) 48 | 49 | 50 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/velocity_divergence.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from . import CellType 3 | 4 | def velocityDivergence(U, flags): 5 | r""" Calculates the velocity divergence (with boundary cond modifications). This is 6 | essentially a replica of makeRhs in Manta and FluidNet. 7 | 8 | Arguments: 9 | U (Tensor): input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 10 | flags (Tensor): input occupancy grid 11 | Output: 12 | UDiv (Tensor) : output divergence (scalar field). 13 | """ 14 | # Check sizes 15 | divergence = torch.zeros_like(flags).type(U.type()) 16 | assert (U.dim() == 5 and flags.dim() == 5 and divergence.dim() == 5), \ 17 | "Dimension mismatch" 18 | assert flags.size(1) == 1, "flags is not scalar" 19 | bsz = flags.size(0) 20 | d = flags.size(2) 21 | h = flags.size(3) 22 | w = flags.size(4) 23 | 24 | z = 2 25 | y = 3 26 | x = 4 27 | 28 | is_3d = (U.size(1) == 3) 29 | if (not is_3d): 30 | assert d == 1, "2D velocity field but zdepth > 1" 31 | assert (U.size(1) == 2), "2D velocity field must have only 2 channels" 32 | 33 | assert (U.size(0) == bsz and U.size(2) == d and \ 34 | U.size(3) == h and U.size(4) == w), "Size mismatch" 35 | 36 | assert (U.is_contiguous() and flags.is_contiguous() and \ 37 | divergence.is_contiguous()), "Input is not contiguous" 38 | 39 | #Uijk : Velocity in ijk 40 | #Uijk_p : Velocity in (i+1),(j+1),(k+1) 41 | 42 | # We call this only on fluid cells. Non-fluid cells have a zero divergence 43 | #isFluid = flags.eq(CellType.TypeFluid) 44 | #noFluid = isFluid.ne(CellType.TypeFluid) 45 | 46 | if (not is_3d): 47 | Uijk = U.narrow(x, 1, w-2).narrow(y, 1, h-2) 48 | Uijk_p = Uijk.clone() 49 | Uijk_p[:,0] = U.narrow(x, 2, w-2).narrow(y, 1, h-2).select(1,0) 50 | Uijk_p[:,1] = U.narrow(x, 1, w-2).narrow(y, 2, h-2).select(1,1) 51 | else: 52 | Uijk = U.narrow(x, 1, w-2).narrow(y, 1, h-2).narrow(z, 1, d-2) 53 | Uijk_p = Uijk.clone() 54 | Uijk_p[:,0] = U.narrow(x, 2, w-2).narrow(y, 1, h-2).narrow(z, 1, d-2).select(1,0) 55 | Uijk_p[:,1] = U.narrow(x, 1, w-2).narrow(y, 2, h-2).narrow(z, 1, d-2).select(1,1) 56 | Uijk_p[:,2] = U.narrow(x, 1, w-2).narrow(y, 1, h-2).narrow(z, 2, d-2).select(1,2) 57 | 58 | # -div = u(i+1,j,k) - u(i,j,k) + 59 | # v(i,j+1,k) - v(i,j,k) + 60 | # w(i,j,k+1) - w(i,j,k) 61 | div = Uijk.select(1,0) - Uijk_p.select(1,0) + \ 62 | Uijk.select(1,1) - Uijk_p.select(1,1) 63 | 64 | if (is_3d): 65 | div += Uijk.select(1,2) - Uijk_p.select(1,2) 66 | if (not is_3d): 67 | divergence[:,:,:,1:(h-1),1:(w-1)] = div.view(bsz, 1, d, h-2, w-2) 68 | else: 69 | divergence[:,:,1:(d-1),1:(h-1),1:(w-1)] = div.view(bsz, 1, d-2, h-2, w-2) 70 | 71 | #Set div to 0 in obstacles 72 | mask_obst = flags.eq(CellType.TypeObstacle) 73 | divergence.masked_fill_(mask_obst, 0) 74 | return divergence 75 | 76 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/velocity_update.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import CellType 4 | 5 | 6 | def velocityUpdate(pressure, U, flags): 7 | 8 | r""" Calculate the pressure gradient and subtract it into (i.e. calculate 9 | U' = U - grad(p)). Some care must be taken with handling boundary conditions. 10 | This function mimics correctVelocity in Manta. 11 | Velocity update is done IN-PLACE. 12 | 13 | Arguments: 14 | p (Tensor): scalar pressure field. 15 | U (Tensor): velocity field (size(2) can be 2 or 3, indicating 2D / 3D) 16 | flags (Tensor): input occupancy grid 17 | """ 18 | # Check arguments. 19 | assert U.dim() == 5 and flags.dim() == 5 and pressure.dim() == 5, \ 20 | "Dimension mismatch" 21 | assert flags.size(1) == 1, "flags is not scalar" 22 | b = flags.size(0) 23 | d = flags.size(2) 24 | h = flags.size(3) 25 | w = flags.size(4) 26 | 27 | is3D = (U.size(1) == 3) 28 | if not is3D: 29 | assert d == 1, "d > 1 for a 2D domain" 30 | assert U.size(4) == w, "2D velocity field must have only 2 channels" 31 | 32 | assert U.size(0) == b and U.size(2) == d and U.size(3) == h \ 33 | and U.size(4) == w, "size mismatch" 34 | assert pressure.is_same_size(flags), "size mismatch" 35 | assert U.is_contiguous() and flags.is_contiguous() and \ 36 | pressure.is_contiguous(), "Input is not contiguous" 37 | 38 | # First, we build the mask for detecting fluid cells. Borders are left untouched. 39 | # mask_fluid Fluid cells. 40 | # mask_fluid_i Fluid cells with (i-1) neighbour also a fluid. 41 | # mask_fluid_j Fluid cells with (j-1) neighbour also a fluid. 42 | # mask_fluid_k FLuid cells with (k-1) neighbour also a fluid. 43 | 44 | # Second, we detect obstacle cells 45 | # See Bridson p44 for algorithm and boundaries treatment. 46 | 47 | if not is3D: 48 | # Current cell is fluid 49 | mask_fluid = flags.narrow(4, 1, w-2).narrow(3, 1, h-2).eq(CellType.TypeFluid) 50 | # Current is fluid and neighbour to left or down are fluid 51 | mask_fluid_i = mask_fluid.__and__ \ 52 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).eq(CellType.TypeFluid)) 53 | mask_fluid_j = mask_fluid.__and__ \ 54 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).eq(CellType.TypeFluid)) 55 | # Current cell is fluid and neighbours to left or down are obstacle 56 | mask_fluid_obstacle_im1 = mask_fluid.__and__ \ 57 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).eq(CellType.TypeEmpty)) 58 | mask_fluid_obstacle_jm1 = mask_fluid.__and__ \ 59 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).eq(CellType.TypeEmpty)) 60 | # Current cell is obstacle and not outflow 61 | mask_obstacle = flags.narrow(4, 1, w-2).narrow(3, 1, h-2) \ 62 | .eq(CellType.TypeEmpty).__and__ \ 63 | (flags.narrow(4, 1, w-2).narrow(3, 1, h-2) \ 64 | .ne(CellType.TypeOutflow)) 65 | # Current cell is obstacle and neighbours to left or down are fluid 66 | mask_obstacle_fluid_im1 = mask_obstacle.__and__ \ 67 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).eq(CellType.TypeFluid)) 68 | mask_obstacle_fluid_jm1 = mask_obstacle.__and__ \ 69 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).eq(CellType.TypeFluid)) 70 | # Current cell is obstacle and neighbours to left or down are not fluid 71 | mask_no_fluid_im1 = mask_obstacle.__and__ \ 72 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).eq(CellType.TypeEmpty)) 73 | mask_no_fluid_jm1 = mask_obstacle.__and__ \ 74 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).eq(CellType.TypeEmpty)) 75 | 76 | else: 77 | # TODO: implement 3D bcs well. 78 | # TODO: add outlfow (change in advection required) 79 | mask_fluid = flags.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2).eq(CellType.TypeFluid) 80 | mask_fluid_i = mask_fluid.__and__ \ 81 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2).eq(CellType.TypeFluid)) 82 | mask_fluid_j = mask_fluid.__and__ \ 83 | (flags.narrow(4, 1, w-1).narrow(3, 0, h-1).narrow(2, 1, d-2).eq(CellType.TypeFluid)) 84 | mask_fluid_k = mask_fluid.__and__ \ 85 | (flags.narrow(4, 1, w-1).narrow(3, 1, h-1).narrow(2, 0, d-2).eq(CellType.TypeFluid)) 86 | 87 | # Cast into float or double tensor and cat into a single mask along chan. 88 | mask_fluid_i_f = mask_fluid_i.type(U.type()) 89 | mask_fluid_j_f = mask_fluid_j.type(U.type()) 90 | 91 | mask_fluid_obstacle_i_f = mask_fluid_obstacle_im1.type(U.type()) 92 | mask_fluid_obstacle_j_f = mask_fluid_obstacle_jm1.type(U.type()) 93 | 94 | mask_obstacle_fluid_i_f = mask_obstacle_fluid_im1.type(U.type()) 95 | mask_obstacle_fluid_j_f = mask_obstacle_fluid_jm1.type(U.type()) 96 | 97 | mask_no_fluid_i_f = mask_no_fluid_im1.type(U.type()) 98 | mask_no_fluid_j_f = mask_no_fluid_jm1.type(U.type()) 99 | 100 | if is3D: 101 | mask_fluid_k_f = mask_fluid_k.type(U.type()) 102 | 103 | if not is3D: 104 | mask_fluid = torch.cat((mask_fluid_i_f, mask_fluid_j_f), 1).contiguous() 105 | mask_fluid_obstacle = torch.cat((mask_fluid_obstacle_i_f, mask_fluid_obstacle_j_f), 1).contiguous() 106 | mask_obstacle_fluid = torch.cat((mask_obstacle_fluid_i_f, mask_obstacle_fluid_j_f), 1).contiguous() 107 | mask_no_fluid = torch.cat((mask_no_fluid_i_f, mask_no_fluid_j_f), 1).contiguous() 108 | else: 109 | mask_fluid = torch.cat((mask_fluid_i_f, mask_fluid_j_f, mask_fluid_k_f), 1).contiguous() 110 | 111 | # pressure tensor. 112 | #Pijk Pressure at (i,j,k) in 3 channels (2 for 2D). 113 | #Pijk_m Pressure at chan 0: (i-1, j, k) 114 | # chan 1: (i, j-1, k) 115 | # chan 2: (i, j, k-1) 116 | 117 | if not is3D: 118 | Pijk = pressure.narrow(4, 1, w-2).narrow(3, 1, h-2) 119 | Pijk = Pijk.clone().expand(b, 2, d, h-2, w-2) 120 | Pijk_m = Pijk.clone().expand(b, 2, d, h-2, w-2) 121 | Pijk_m[:,0] = pressure.narrow(4, 0, w-2).narrow(3, 1, h-2).squeeze(1) 122 | Pijk_m[:,1] = pressure.narrow(4, 1, w-2).narrow(3, 0, h-2).squeeze(1) 123 | else: 124 | Pijk = pressure.narrow(4, 1, w-1).narrow(3, 1, h-1).narrow(2, 1, d-2) 125 | Pijk = Pijk.clone().expand(b, 3, d-2, h-1, w-1) 126 | Pijk_m = Pijk.clone().expand(b, 3, d-2, h-1, w-1) 127 | Pijk_m[:,0] = pressure.narrow(4, 0, w-1).narrow(3, 1, h-1).narrow(2, 1, d-2).squeeze(1) 128 | Pijk_m[:,1] = pressure.narrow(4, 1, w-1).narrow(3, 0, h-1).narrow(2, 1, d-2).squeeze(1) 129 | Pijk_m[:,2] = pressure.narrow(4, 1, w-1).narrow(3, 1, h-1).narrow(2, 0, d-2).squeeze(1) 130 | 131 | # grad(p) = [[ p(i,j,k) - p(i-1,j,k) ] 132 | # [ p(i,j,k) - p(i,j-1,k) ] 133 | # [ p(i,j,k) - p(i,j,k-1) ]] 134 | if not is3D: 135 | # Three cases: 136 | # 1) Cell is fluid and left neighbour is fluid: 137 | # u = u - grad(p) 138 | # 2) Cell is fluid and left neighbour is obstacle 139 | # u = u - p(i,j) 140 | # 3) Cell is obstacle and left neighbour is fluid 141 | # u = u + p(i-1,j) 142 | 143 | U[:,:,:,1:(h-1),1:(w-1)] = (mask_fluid * \ 144 | (U.narrow(4, 1, w-2).narrow(3, 1, h-2) - (Pijk - Pijk_m)) + \ 145 | mask_fluid_obstacle * \ 146 | (U.narrow(4, 1, w-2).narrow(3, 1, h-2) - Pijk) + \ 147 | mask_obstacle_fluid * \ 148 | (U.narrow(4, 1, w-2).narrow(3, 1, h-2) + Pijk_m) + \ 149 | mask_no_fluid * (0)) 150 | #print('******************************************************') 151 | 152 | #print('masks') 153 | #print('* * * *') 154 | #print(mask_fluid) 155 | #print('* * * *') 156 | #print(mask_fluid_obstacle) 157 | #print(mask_obstacle_fluid) 158 | #print(mask_no_fluid) 159 | #print('******************************************************') 160 | else: 161 | U[:,:,1:(d-1),1:(h-1),1:(w-1)] = mask * \ 162 | (U.narrow(4, 1, w-1).narrow(3, 1, h-1).narrow(2, 1, d-2) - (Pijk - Pijk_m)) 163 | 164 | -------------------------------------------------------------------------------- /pytorch/lib/fluid/viscosity.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from . import CellType 4 | 5 | # ******WARNING********: ONLY IN 2D 6 | 7 | def addViscosity(dt, U, flags, viscosity): 8 | # Check arguments. 9 | r"""Adds viscosity to velocity field. 10 | Velocity update is done IN-PLACE. 11 | 12 | Arguments: 13 | p (Tensor): scalar pressure field. 14 | U (Tensor): velocity field (size(2) can be 2 or 3, indicating 2D / 3D) 15 | flags (Tensor): input occupancy grid 16 | """ 17 | assert U.dim() == 5 and flags.dim() == 5, "Dimension mismatch" 18 | assert flags.size(1) == 1, "flags is not scalar" 19 | b = flags.size(0) 20 | d = flags.size(2) 21 | h = flags.size(3) 22 | w = flags.size(4) 23 | 24 | is3D = (U.size(1) == 3) 25 | if not is3D: 26 | assert d == 1, "d > 1 for a 2D domain" 27 | assert U.size(4) == w, "2D velocity field must have only 2 channels" 28 | 29 | assert U.size(0) == b and U.size(2) == d and U.size(3) == h \ 30 | and U.size(4) == w, "size mismatch" 31 | assert U.is_contiguous() and flags.is_contiguous(), "Input is not contiguous" 32 | 33 | # First, we build the mask for detecting fluid cells. Borders are left untouched. 34 | # mask_fluid Fluid cells. 35 | # mask_fluid_i Fluid cells with (i-1) neighbour also a fluid. 36 | # mask_fluid_j Fluid cells with (j-1) neighbour also a fluid. 37 | # mask_fluid_k Fluid cells with (k-1) neighbour also a fluid. 38 | 39 | if not is3D: 40 | mask_fluid = flags.narrow(4, 1, w-2).narrow(3, 1, h-2).eq(CellType.TypeFluid) 41 | mask_fluid_i = mask_fluid.__and__ \ 42 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).eq(CellType.TypeFluid)) 43 | mask_fluid_j = mask_fluid.__and__ \ 44 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).eq(CellType.TypeFluid)) 45 | 46 | # Cast into float or double tensor and cat into a single mask along chan. 47 | mask_fluid_i_f = mask_fluid_i.type(U.type()) 48 | mask_fluid_j_f = mask_fluid_j.type(U.type()) 49 | if is3D: 50 | mask_fluid_k_f = mask_fluid_k.type(U.type()) 51 | 52 | if not is3D: 53 | mask = torch.cat((mask_fluid_i_f, mask_fluid_j_f), 1).contiguous() 54 | else: 55 | mask = torch.cat((mask_fluid_i_f, mask_fluid_j_f, mask_fluid_k_f), 1).contiguous() 56 | 57 | # Update the velocity, to the viscous velocity field. 58 | # u^v(i,j) = u^n(i,j) 59 | # + dt*nu*[u^n(i+1,j) + u^n(i,j+1) + u^n(i-1,j) + u^n(i,j-1) 60 | # -4u^n(i,j) ] 61 | if not is3D: 62 | U[:,:,:,1:(h-1),1:(w-1)] = mask * (\ 63 | U.narrow(4, 1, w-2).narrow(3, 1, h-2) + 64 | dt * viscosity *( 65 | U.narrow(4, 2, w-2).narrow(3, 1, h-2) + \ 66 | U.narrow(4, 1, w-2).narrow(3, 2, h-2) + \ 67 | U.narrow(4, 0, w-2).narrow(3, 1, h-2) + \ 68 | U.narrow(4, 0, w-2).narrow(3, 0, h-2) - \ 69 | (4*U.narrow(4, 1, w-2).narrow(3, 1, h-2) )\ 70 | ) ) 71 | -------------------------------------------------------------------------------- /pytorch/lib/load_manta_data.py: -------------------------------------------------------------------------------- 1 | import struct 2 | import torch 3 | 4 | def loadMantaFile(fname): 5 | with open(fname, 'rb') as f: 6 | fhead = struct.unpack('i'*5, f.read(4*5)) 7 | nx = fhead[1] 8 | ny = fhead[2] 9 | nz = fhead[3] 10 | is3D = (fhead[4] == 1) 11 | 12 | numel = nx*ny*nz 13 | 14 | # read moves cursor to the end of previous read 15 | # we don't need to move the cursor using seek 16 | ld_array = struct.unpack('f'*3*numel, f.read(4*3*numel)) 17 | Ux = torch.FloatTensor(ld_array[:numel]) 18 | Uy = torch.FloatTensor(ld_array[numel:(numel*2)]) 19 | p = torch.FloatTensor(ld_array[2*numel:(numel*3)]) 20 | if (is3D): 21 | Uz = torch.FloatTensor(struct.unpack('f'*numel, f.read(4*numel))) 22 | 23 | flags = torch.IntTensor(struct.unpack('i'*numel, f.read(4*numel))).float() 24 | density = torch.FloatTensor(struct.unpack('f'*numel, f.read(4*numel))) 25 | 26 | # We ALWAYS deal with 5D tensor to make things easier. 27 | # All tensor are always nbatch x nchan x nz x ny x nx. 28 | Ux.resize_(1,1,nz,ny,nx) 29 | Uy.resize_(1,1,nz,ny,nx) 30 | if (is3D): 31 | Uz.resize_(1,1,nz,ny,nx) 32 | p.resize_(1,1,nz,ny,nx) 33 | flags.resize_(1,1,nz,ny,nx) 34 | density.resize_(1,1,nz,ny,nx) 35 | 36 | if (is3D): 37 | U = torch.cat((Ux, Uy, Uz), 1).contiguous() 38 | else: 39 | U = torch.cat((Ux, Uy), 1).contiguous() 40 | 41 | return p, U, flags, density, is3D 42 | -------------------------------------------------------------------------------- /pytorch/lib/model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from lib import fluid, MultiScaleNet 6 | from math import inf 7 | 8 | class _ScaleNet(nn.Module): 9 | def __init__(self, mconf): 10 | super(_ScaleNet, self).__init__() 11 | self.mconf = mconf 12 | 13 | def forward(self, x): 14 | bsz = x.size(0) 15 | # Rehaspe form (b x chan x d x h x w) to (b x -1) 16 | y = x.view(bsz, -1) 17 | # Calculate std using Bessel's correction (correction with n/n-1) 18 | std = torch.std(y, dim=1, keepdim=True) # output is size (b x 1) 19 | scale = torch.clamp(std, \ 20 | self.mconf['normalizeInputThreshold'] , inf) 21 | scale = scale.view(bsz, 1, 1, 1, 1) 22 | 23 | return scale 24 | 25 | class _HiddenConvBlock(nn.Module): 26 | def __init__(self, dropout=True): 27 | super(_HiddenConvBlock, self).__init__() 28 | layers = [ 29 | nn.Conv2d(16, 16, kernel_size=3, padding = 1), 30 | nn.ReLU(inplace=True), 31 | nn.Conv2d(16, 16, kernel_size=3, padding = 1), 32 | nn.ReLU(), 33 | ] 34 | if dropout: 35 | layers.append(nn.Dropout()) 36 | self.block = nn.Sequential(*layers) 37 | 38 | def forward(self, x): 39 | return self.block(x) 40 | 41 | class FluidNet(nn.Module): 42 | # For now, only 2D model. Add 2D/3D option. Only known from data! 43 | # Also, build model with MSE of pressure as loss func, therefore input is velocity 44 | # and output is pressure, to be compared to target pressure. 45 | def __init__(self, mconf, dropout=True): 46 | super(FluidNet, self).__init__() 47 | 48 | 49 | self.dropout = dropout 50 | self.mconf = mconf 51 | self.inDims = mconf['inputDim'] 52 | self.is3D = mconf['is3D'] 53 | 54 | self.scale = _ScaleNet(self.mconf) 55 | # Input channels = 3 (inDims, flags) 56 | # We add padding to make sure that Win = Wout and Hin = Hout with ker_size=3 57 | self.conv1 = torch.nn.Conv2d(self.inDims, 16, kernel_size=3, padding=1) 58 | 59 | self.modDown1 = torch.nn.AvgPool2d(kernel_size=2) 60 | self.modDown2 = torch.nn.AvgPool2d(kernel_size=4) 61 | 62 | self.convBank = _HiddenConvBlock(dropout=False) 63 | 64 | #self.deconv1 = torch.nn.ConvTranspose2d(16, 16, kernel_size=2, stride=2) 65 | #self.deconv2 = torch.nn.ConvTranspose2d(16, 16, kernel_size=4, stride=4) 66 | 67 | self.conv2 = torch.nn.Conv2d(16, 16, kernel_size=1) 68 | self.conv3 = torch.nn.Conv2d(16, 8, kernel_size=1) 69 | 70 | # Output channels = 1 (pressure) 71 | self.convOut = torch.nn.Conv2d(8, 1, kernel_size=1) 72 | 73 | # MultiScaleNet 74 | self.multiScale = MultiScaleNet(self.inDims) 75 | 76 | def forward(self, input_): 77 | 78 | # data indexes | | 79 | # (dim 1) | 2D | 3D 80 | # ---------------------------------------- 81 | # DATA: 82 | # pDiv | 0 | 0 83 | # UDiv | 1:3 | 1:4 84 | # flags | 3 | 4 85 | # densityDiv | 4 | 5 86 | # TARGET: 87 | # p | 0 | 0 88 | # U | 1:3 | 1:4 89 | # density | 3 | 4 90 | 91 | # For now, we work ONLY in 2d 92 | 93 | assert self.is3D == False, 'Input can only be 2D' 94 | 95 | assert self.mconf['inputChannels']['pDiv'] or \ 96 | self.mconf['inputChannels']['UDiv'] or \ 97 | self.mconf['inputChannels']['div'], 'Choose at least one field (U, div or p).' 98 | 99 | pDiv = None 100 | UDiv = None 101 | div = None 102 | 103 | # Flags are always loaded 104 | if self.is3D: 105 | flags = input_[:,4].unsqueeze(1) 106 | else: 107 | flags = input_[:,3].unsqueeze(1).contiguous() 108 | 109 | if (self.mconf['inputChannels']['pDiv'] or (self.mconf['normalizeInput'] \ 110 | and self.mconf['normalizeInputChan'] == 'pDiv')): 111 | pDiv = input_[:,0].unsqueeze(1).contiguous() 112 | 113 | if (self.mconf['inputChannels']['UDiv'] or self.mconf['inputChannels']['div'] \ 114 | or (self.mconf['normalizeInput'] \ 115 | and self.mconf['normalizeInputChan'] == 'UDiv')): 116 | if self.is3D: 117 | UDiv = input_[:,1:4].contiguous() 118 | else: 119 | UDiv = input_[:,1:3].contiguous() 120 | 121 | # Apply setWallBcs to zero out obstacles velocities on the boundary 122 | # XXX: This seems wrong to put here. 123 | #UDiv = fluid.setWallBcs(UDiv, flags) 124 | 125 | if self.mconf['inputChannels']['div']: 126 | div = fluid.velocityDivergence(UDiv, flags) 127 | 128 | # Apply scale to input 129 | if self.mconf['normalizeInput']: 130 | if self.mconf['normalizeInputChan'] == 'UDiv': 131 | s = self.scale(UDiv) 132 | elif self.mconf['normalizeInputChan'] == 'pDiv': 133 | s = self.scale(pDiv) 134 | elif self.mconf['normalizeInputChan'] == 'div': 135 | s = self.scale(div) 136 | else: 137 | raise Exception('Incorrect normalize input channel.') 138 | 139 | if pDiv is not None: 140 | pDiv = torch.div(pDiv, s) 141 | if UDiv is not None: 142 | UDiv = torch.div(UDiv, s) 143 | if div is not None: 144 | div = torch.div(div, s) 145 | 146 | x = torch.FloatTensor(input_.size(0), \ 147 | self.inDims, \ 148 | input_.size(2), \ 149 | input_.size(3), \ 150 | input_.size(4)).type_as(input_) 151 | 152 | chan = 0 153 | if self.mconf['inputChannels']['pDiv']: 154 | x[:, chan] = pDiv[:,0] 155 | chan += 1 156 | elif self.mconf['inputChannels']['UDiv']: 157 | if self.is3D: 158 | x[:,chan:(chan+3)] = UDiv 159 | chan += 3 160 | else: 161 | x[:,chan:(chan+2)] = UDiv 162 | chan += 2 163 | elif self.mconf['inputChannels']['div']: 164 | x[:, chan] = div[:,0] 165 | chan += 1 166 | 167 | # FlagsToOccupancy creates a [0,1] grid out of the manta flags 168 | x[:,chan,:,:,:] = fluid.flagsToOccupancy(flags).squeeze(1) 169 | 170 | if not self.is3D: 171 | # Squeeze unary dimension as we are in 2D 172 | x = torch.squeeze(x,2) 173 | 174 | if self.mconf['model'] == 'ScaleNet': 175 | p = self.multiScale(x) 176 | 177 | else: 178 | # Inital layers 179 | x = F.relu(self.conv1(x)) 180 | 181 | # We divide the network in 3 banks, applying average pooling 182 | x1 = self.modDown1(x) 183 | x2 = self.modDown2(x) 184 | 185 | # Process every bank in parallel 186 | x0 = self.convBank(x) 187 | x1 = self.convBank(x1) 188 | x2 = self.convBank(x2) 189 | 190 | # Upsample banks 1 and 2 to bank 0 size and accumulate inputs 191 | 192 | #x1 = self.upscale1(x1) 193 | #x2 = self.upscale2(x2) 194 | 195 | x1 = F.interpolate(x1, scale_factor=2) 196 | x2 = F.interpolate(x2, scale_factor=4) 197 | #x1 = self.deconv1(x1) 198 | #x2 = self.deconv2(x2) 199 | 200 | #x = torch.cat((x0, x1, x2), dim=1) 201 | x = x0 + x1 + x2 202 | 203 | # Apply last 2 convolutions 204 | x = F.relu(self.conv2(x)) 205 | x = F.relu(self.conv2(x)) 206 | x = F.relu(self.conv3(x)) 207 | 208 | # Output pressure (1 chan) 209 | p = self.convOut(x) 210 | 211 | 212 | # Add back the unary dimension 213 | if not self.is3D: 214 | p = torch.unsqueeze(p, 2) 215 | 216 | # Correct U = UDiv - grad(p) 217 | # flags is the one with Manta's values, not occupancy in [0,1] 218 | fluid.velocityUpdate(pressure=p, U=UDiv, flags=flags) 219 | 220 | # We now UNDO the scale factor we applied on the input. 221 | if self.mconf['normalizeInput']: 222 | p = torch.mul(p,s) # Applies p' = *= scale 223 | UDiv = torch.mul(UDiv,s) 224 | 225 | # Set BCs after velocity update. 226 | UDiv = fluid.setWallBcs(UDiv, flags) 227 | return p, UDiv 228 | 229 | 230 | -------------------------------------------------------------------------------- /pytorch/lib/multi_scale_net.py: -------------------------------------------------------------------------------- 1 | """ 2 | 3 level multiscale network 3 | 4 | Inputs are shape (batch, channels, height, width) 5 | Outputs are shape (batch,1, height, width) 6 | 7 | The number of input (data) channels is selected when the model is created. 8 | the number of output (target) channels is fixed at 1, although this could be changed in the future. 9 | 10 | The data can be any size (i.e. height and width), although for best results the height and width should 11 | be divisble by four. 12 | 13 | The model can be trained on data of a given size (H and W) and then used on data of any other size, 14 | although the best results so far have been obtained with test data of similar size to the training data 15 | 16 | """ 17 | import torch 18 | import torch.nn as nn 19 | import torch.nn.functional as F 20 | 21 | class _ConvBlock1(nn.Module): 22 | """ 23 | First block - quarter scale. 24 | Four Conv2d layers, all with kernel_size 3 and padding of 1 (padding ensures output size is same as input size) 25 | Optional dropout before final Conv2d layer 26 | ReLU after first two Conv2d layers, not after last two - predictions can be +ve or -ve 27 | """ 28 | def __init__(self, in_channels, mid1_channels, mid2_channels, out_channels,dropout=True): 29 | super(_ConvBlock1, self).__init__() 30 | layers = [ 31 | nn.Conv2d(in_channels, mid1_channels, kernel_size=3, padding = 1), 32 | nn.ReLU(inplace=True), 33 | nn.Conv2d(mid1_channels, mid2_channels, kernel_size=3, padding = 1), 34 | nn.ReLU(), 35 | nn.Conv2d(mid2_channels,mid1_channels,kernel_size = 3, padding = 1), 36 | ] 37 | if dropout: 38 | layers.append(nn.Dropout()) 39 | layers.append(nn.Conv2d(mid1_channels, out_channels, kernel_size = 3, padding = 1)) 40 | self.encode = nn.Sequential(*layers) 41 | 42 | def forward(self, x): 43 | return self.encode(x) 44 | 45 | class _ConvBlock2(nn.Module): 46 | """ 47 | Second block - half scale. 48 | Six Conv2d layers. First one kernel size 5, padding 2, remainder kernel size 3 padding 1. 49 | Optional dropout before final Conv2d layer 50 | ReLU after first four Conv2d layers, not after last two - predictions can be +ve or -ve 51 | """ 52 | def __init__(self, in_channels, mid1_channels, mid2_channels,mid3_channels, out_channels,dropout=True): 53 | super(_ConvBlock2, self).__init__() 54 | layers = [ 55 | nn.Conv2d(in_channels, mid1_channels, kernel_size=5, padding = 2), 56 | nn.ReLU(inplace=True), 57 | nn.Conv2d(mid1_channels, mid2_channels, kernel_size=3, padding = 1), 58 | nn.ReLU(), 59 | nn.Conv2d(mid2_channels,mid3_channels,kernel_size = 3, padding = 1), 60 | nn.ReLU(), 61 | nn.Conv2d(mid3_channels,mid2_channels,kernel_size = 3, padding = 1), 62 | nn.ReLU(), 63 | nn.Conv2d(mid2_channels,mid1_channels,kernel_size = 3, padding = 1), 64 | ] 65 | if dropout: 66 | layers.append(nn.Dropout()) 67 | layers.append(nn.Conv2d(mid1_channels, out_channels, kernel_size = 3, padding = 1)) 68 | self.encode = nn.Sequential(*layers) 69 | 70 | def forward(self, x): 71 | return self.encode(x) 72 | 73 | class _ConvBlock3(nn.Module): 74 | """ 75 | Third block - full scale. 76 | Six Conv2d layers. First and last kernel size 5, padding 2, remainder kernel size 3 padding 1. 77 | Optional dropout before final Conv2d layer 78 | ReLU after first four Conv2d layers, not after last two - predictions can be +ve or -ve 79 | """ 80 | def __init__(self, in_channels, mid1_channels, mid2_channels,mid3_channels, out_channels,dropout=True): 81 | super(_ConvBlock3, self).__init__() 82 | layers = [ 83 | nn.Conv2d(in_channels, mid1_channels, kernel_size=5, padding = 2), 84 | nn.ReLU(inplace=True), 85 | nn.Conv2d(mid1_channels, mid2_channels, kernel_size=3, padding = 1), 86 | nn.ReLU(), 87 | nn.Conv2d(mid2_channels,mid3_channels,kernel_size = 3, padding = 1), 88 | nn.ReLU(), 89 | nn.Conv2d(mid3_channels,mid2_channels,kernel_size = 3, padding = 1), 90 | nn.ReLU(), 91 | nn.Conv2d(mid2_channels,mid1_channels,kernel_size = 3, padding = 1), 92 | ] 93 | if dropout: 94 | layers.append(nn.Dropout()) 95 | layers.append(nn.Conv2d(mid1_channels, out_channels, kernel_size = 5, padding = 2)) 96 | self.encode = nn.Sequential(*layers) 97 | 98 | def forward(self, x): 99 | return self.encode(x) 100 | 101 | class MultiScaleNet(nn.Module): 102 | """ 103 | Define the network. Only input when called is number of data (input) channels. 104 | -Downsample input to quarter scale and use ConvBlock1. 105 | -Upsample output of ConvBlock1 to half scale. 106 | -Downsample input to half scale, concat to output of ConvBlock1; use ConvBLock2. 107 | -Upsample output of ConvBlock2 to full scale. 108 | -Concat input to output of ConvBlock2, use ConvBlock3. Output of ConvBlock3 has 8 channels 109 | -Use final Conv2d layer with kernel size of 1 to go from 8 channels to 1 output channel. 110 | """ 111 | def __init__(self,data_channels): 112 | super(MultiScaleNet, self).__init__() 113 | self.convN_4 = _ConvBlock1(data_channels, 32,64,1) 114 | self.convN_2 = _ConvBlock2(data_channels+1, 32,64,128,1) 115 | self.convN_1 = _ConvBlock3(data_channels+1, 32,64,128,8) 116 | self.final = nn.Conv2d(8,1, kernel_size = 1) 117 | 118 | def forward(self, x): 119 | quarter_size = [int(i*0.25) for i in list(x.size()[2:])] 120 | half_size = [int(i*0.5) for i in list(x.size()[2:])] 121 | convN_4out = self.convN_4(F.upsample(x,(quarter_size),mode = 'bilinear')) 122 | convN_2out = self.convN_2( torch.cat((F.upsample(x,(half_size),mode = 'bilinear'), 123 | F.upsample(convN_4out,(half_size),mode = 'bilinear')),dim = 1) ) 124 | convN_1out = self.convN_1( torch.cat((F.upsample(x,(x.size()[2:]),mode = 'bilinear'), 125 | F.upsample(convN_2out,(x.size()[2:]),mode = 'bilinear')),dim = 1) ) 126 | final_out = self.final(convN_1out) 127 | return final_out 128 | -------------------------------------------------------------------------------- /pytorch/lib/simulate.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import lib.fluid as fluid 3 | 4 | def setConstVals(batch_dict, p, U, flags, density): 5 | # apply external BCs. 6 | # batch_dict at output = {p, UDiv, flags, density, UBC, 7 | # UBCInvMask, densityBC, densityBCInvMask} 8 | 9 | #if 'cylinder' in batch_dict: 10 | # # Zero out the U values on the BCs. 11 | # U.mul_(batch_dict['InvInletMask']) 12 | # # Add back the values we want to specify. 13 | # U.add_(batch_dict['UInlet']) 14 | # batch_dict['U'] = U.clone() 15 | 16 | if ('UBCInvMask' in batch_dict) and ('UBC' in batch_dict): 17 | # Zero out the U values on the BCs. 18 | U.mul_(batch_dict['UBCInvMask']) 19 | # Add back the values we want to specify. 20 | U.add_(batch_dict['UBC']) 21 | batch_dict['U'] = U.clone() 22 | 23 | if ('densityBCInvMask' in batch_dict) and ('densityBC' in batch_dict): 24 | density.mul_(batch_dict['densityBCInvMask']) 25 | density.add_(batch_dict['densityBC']) 26 | batch_dict['density'] = density.clone() 27 | 28 | def simulate(mconf, batch_dict, net, sim_method, output_div=False): 29 | r"""Top level simulation loop. 30 | 31 | Arguments: 32 | mconf (dict): Model configuration dictionnary. 33 | batch_dict (dict): Dictionnary of torch Tensors. 34 | Keys must be 'U', 'flags', 'p', 'density'. 35 | Simulations are done INPLACE. 36 | net (nn.Module): convNet model. 37 | sim_method (string): Options are 'convnet' and 'jacobi' 38 | output_div (bool, optional): returns just before solving for pressure. 39 | i.e. leave the state as UDiv and pDiv (before substracting divergence) 40 | 41 | """ 42 | cuda = torch.device('cuda') 43 | assert sim_method == 'convnet' or sim_method == 'jacobi', 'Simulation method \ 44 | not supported. Choose either convnet or jacobi.' 45 | 46 | dt = float(mconf['dt']) 47 | maccormackStrength = mconf['maccormackStrength'] 48 | sampleOutsideFluid = mconf['sampleOutsideFluid'] 49 | 50 | buoyancyScale = mconf['buoyancyScale'] 51 | gravityScale = mconf['gravityScale'] 52 | 53 | viscosity = mconf['viscosity'] 54 | assert viscosity >= 0, 'Viscosity must be positive' 55 | 56 | # Get p, U, flags and density from batch. 57 | p = batch_dict['p'] 58 | U = batch_dict['U'] 59 | 60 | flags = batch_dict['flags'] 61 | stick = False 62 | if 'flags_stick' in batch_dict: 63 | stick = True 64 | flags_stick = batch_dict['flags_stick'] 65 | 66 | # If viscous model, add viscosity 67 | if (viscosity > 0): 68 | orig = U.clone() 69 | fluid.addViscosity(dt, orig, flags, viscosity) 70 | 71 | if 'density' in batch_dict: 72 | density = batch_dict['density'] 73 | 74 | # First advect all scalar fields. 75 | density = fluid.advectScalar(dt, density, U, flags, \ 76 | method="maccormackFluidNet", \ 77 | boundary_width=1, sample_outside_fluid=sampleOutsideFluid, \ 78 | maccormack_strength=maccormackStrength) 79 | if mconf['correctScalar']: 80 | div = fluid.velocityDivergence(U, flags) 81 | fluid.correctScalar(dt, density, div, flags) 82 | else: 83 | density = torch.zeros_like(flags) 84 | 85 | if viscosity == 0: 86 | # Self-advect velocity if inviscid 87 | U = fluid.advectVelocity(dt=dt, orig=U, U=U, flags=flags, method="maccormackFluidNet", \ 88 | boundary_width=1, maccormack_strength=maccormackStrength) 89 | else: 90 | # Advect viscous velocity field orig by the non-divergent 91 | # velocity field U. 92 | U = fluid.advectVelocity(dt=dt, orig=orig, U=U, flags=flags, method="maccormackFluidNet", \ 93 | boundary_width=1, maccormack_strength=maccormackStrength) 94 | 95 | # Set the manual BCs. 96 | setConstVals(batch_dict, p, U, flags, density) 97 | 98 | if 'density' in batch_dict: 99 | if buoyancyScale > 0: 100 | # Add external forces: buoyancy. 101 | gravity = torch.FloatTensor(3).fill_(0).cuda() 102 | gravity[0] = mconf['gravityVec']['x'] 103 | gravity[1] = mconf['gravityVec']['y'] 104 | gravity[2] = mconf['gravityVec']['z'] 105 | gravity.mul_(-buoyancyScale) 106 | rho_star = mconf['operatingDensity'] 107 | U = fluid.addBuoyancy(U, flags, density, gravity, rho_star, dt) 108 | if gravityScale > 0: 109 | gravity = torch.FloatTensor(3).fill_(0).cuda() 110 | gravity[0] = mconf['gravityVec']['x'] 111 | gravity[1] = mconf['gravityVec']['y'] 112 | gravity[2] = mconf['gravityVec']['z'] 113 | # Add external forces: gravity. 114 | gravity.mul_(-gravityScale) 115 | U = fluid.addGravity(U, flags, gravity, dt) 116 | 117 | if (output_div): 118 | return 119 | 120 | if sim_method != 'convnet': 121 | if 'periodic-x' in mconf and 'periodic-y' in mconf: 122 | U_temp = U.clone() 123 | U = fluid.setWallBcs(U, flags) 124 | if 'periodic-x' in mconf and 'periodic-y' in mconf: 125 | if mconf['periodic-x']: 126 | U[:,1,:,:,1] = U_temp[:,1,:,:,U.size(4)-1] 127 | if mconf['periodic-y']: 128 | U[:,0,:,1] = U_temp[:,0,:,U.size(3)-1] 129 | elif stick: 130 | fluid.setWallBcsStick(U, flags, flags_stick) 131 | 132 | # Set the constant domain values. 133 | setConstVals(batch_dict, p, U, flags, density) 134 | 135 | 136 | if (sim_method == 'convnet'): 137 | # fprop the model to perform the pressure projection and velocity calculation. 138 | # Set wall BCs is performed inside the model, before and after the projection. 139 | # No need to call it again. 140 | net.eval() 141 | data = torch.cat((p, U, flags, density), 1) 142 | p, U = net(data) 143 | 144 | elif (sim_method == 'jacobi'): 145 | div = fluid.velocityDivergence(U, flags) 146 | 147 | is3D = (U.size(2) > 1) 148 | pTol = mconf['pTol'] 149 | maxIter = mconf['jacobiIter'] 150 | 151 | p, residual = fluid.solveLinearSystemJacobi( \ 152 | flags=flags, div=div, is_3d=is3D, p_tol=pTol, \ 153 | max_iter=maxIter) 154 | fluid.velocityUpdate(pressure=p, U=U, flags=flags) 155 | 156 | if sim_method != 'convnet': 157 | if 'periodic-x' in mconf and 'periodic-y' in mconf: 158 | U_temp = U.clone() 159 | U = fluid.setWallBcs(U, flags) 160 | if 'periodic-x' in mconf and 'periodic-y' in mconf: 161 | if mconf['periodic-x']: 162 | U[:,1,:,:,1] = U_temp[:,1,:,:,U.size(4)-1] 163 | if mconf['periodic-y']: 164 | U[:,0,:,1] = U_temp[:,0,:,U.size(3)-1] 165 | elif stick: 166 | fluid.setWallBcsStick(U, flags, flags_stick) 167 | 168 | setConstVals(batch_dict, p, U, flags, density) 169 | batch_dict['U'] = U 170 | batch_dict['density'] = density 171 | batch_dict['p'] = p 172 | -------------------------------------------------------------------------------- /pytorch/lib/util_print.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | from collections import OrderedDict 5 | 6 | def summary(model, input_size): 7 | def register_hook(module): 8 | def hook(module, input, output): 9 | class_name = str(module.__class__).split('.')[-1].split("'")[0] 10 | module_idx = len(summary) 11 | 12 | m_key = '%s-%i' % (class_name, module_idx+1) 13 | summary[m_key] = OrderedDict() 14 | summary[m_key]['input_shape'] = list(input[0].size()) 15 | summary[m_key]['input_shape'][0] = -1 16 | if isinstance(output, (list,tuple)): 17 | summary[m_key]['output_shape'] = [[-1] + list(o.size())[1:] for o in output] 18 | else: 19 | summary[m_key]['output_shape'] = list(output.size()) 20 | summary[m_key]['output_shape'][0] = -1 21 | 22 | params = 0 23 | if hasattr(module, 'weight'): 24 | params += torch.prod(torch.LongTensor(list(module.weight.size()))) 25 | summary[m_key]['trainable'] = module.weight.requires_grad 26 | if hasattr(module, 'bias') and hasattr(module.bias, 'size'): 27 | params += torch.prod(torch.LongTensor(list(module.bias.size()))) 28 | summary[m_key]['nb_params'] = params 29 | 30 | if (not isinstance(module, nn.Sequential) and \ 31 | not isinstance(module, nn.ModuleList) and \ 32 | not (module == model)): 33 | hooks.append(module.register_forward_hook(hook)) 34 | 35 | if torch.cuda.is_available(): 36 | dtype = torch.cuda.FloatTensor 37 | else: 38 | dtype = torch.FloatTensor 39 | 40 | # check if there are multiple inputs to the network 41 | if isinstance(input_size[0], (list, tuple)): 42 | x = [torch.rand(1,*in_size).type(dtype) for in_size in input_size] 43 | else: 44 | x = torch.rand(1,*input_size).type(dtype) 45 | 46 | # print(type(x[0])) 47 | # create properties 48 | summary = OrderedDict() 49 | hooks = [] 50 | # register hook 51 | model.apply(register_hook) 52 | # make a forward pass 53 | print(x.shape) 54 | model(x) 55 | # remove these hooks 56 | for h in hooks: 57 | h.remove() 58 | 59 | print('----------------------------------------------------------------') 60 | line_new = '{:>20} {:>25} {:>15}'.format('Layer (type)', 'Output Shape', 'Param #') 61 | print(line_new) 62 | print('================================================================') 63 | total_params = 0 64 | trainable_params = 0 65 | for layer in summary: 66 | # input_shape, output_shape, trainable, nb_params 67 | line_new = '{:>20} {:>25} {:>15}'.format(layer, str(summary[layer]['output_shape']), summary[layer]['nb_params']) 68 | total_params += summary[layer]['nb_params'] 69 | if 'trainable' in summary[layer]: 70 | if summary[layer]['trainable'] == True: 71 | trainable_params += summary[layer]['nb_params'] 72 | print(line_new) 73 | print('================================================================') 74 | print('Total params: ' + str(total_params)) 75 | print('Trainable params: ' + str(trainable_params)) 76 | print('Non-trainable params: ' + str(total_params - trainable_params)) 77 | print('----------------------------------------------------------------') 78 | # return summary 79 | 80 | -------------------------------------------------------------------------------- /pytorch/plot_5loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import glob 5 | 6 | 7 | # Load numpy arrays 8 | assert (len(sys.argv) == 2), 'Usage: python3 plot_loss.py ' 9 | assert (glob.os.path.exists(sys.argv[1])), 'Directory ' + str(sys.argv[1]) + ' does not exists' 10 | 11 | save_dir = sys.argv[1] 12 | file_train = glob.os.path.join(save_dir, 'train_loss.npy') 13 | file_val = glob.os.path.join(save_dir, 'val_loss.npy') 14 | train_loss_plot = np.load(file_train) 15 | val_loss_plot = np.load(file_val) 16 | 17 | init = 5 18 | 19 | # Plot loss against epochs 20 | plt.plot(train_loss_plot[init:,0], train_loss_plot[init:,1], label = 'Total Training Loss') 21 | plt.plot(train_loss_plot[init:,0], train_loss_plot[init:,2], label = 'L2 pressure Loss') 22 | plt.plot(train_loss_plot[init:,0], train_loss_plot[init:,3], label = 'L2 div Loss') 23 | plt.plot(train_loss_plot[init:,0], train_loss_plot[init:,4], label = 'L1 pressure Loss') 24 | plt.plot(train_loss_plot[init:,0], train_loss_plot[init:,5], label = 'L1 div Loss') 25 | plt.plot(train_loss_plot[init:,0], train_loss_plot[init:,6], label = 'Long Term div Loss') 26 | plt.plot(val_loss_plot[init:,0], val_loss_plot[init:,1], label = 'Validation Loss') 27 | plt.legend() 28 | plt.xlabel('Epoch') 29 | plt.ylabel('Loss') 30 | plt.show(block=True) 31 | 32 | -------------------------------------------------------------------------------- /pytorch/plot_loss.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | import glob 5 | 6 | 7 | # Load numpy arrays 8 | assert (len(sys.argv) == 2), 'Usage: python3 plot_loss.py ' 9 | assert (glob.os.path.exists(sys.argv[1])), 'Directory ' + str(sys.argv[1]) + ' does not exists' 10 | 11 | save_dir = sys.argv[1] 12 | file_train = glob.os.path.join(save_dir, 'train_loss.npy') 13 | file_val = glob.os.path.join(save_dir, 'val_loss.npy') 14 | train_loss_plot = np.load(file_train) 15 | val_loss_plot = np.load(file_val) 16 | 17 | x_start = 0 18 | # Plot loss against epochs 19 | plt.plot(train_loss_plot[x_start:,0], train_loss_plot[x_start:,1], label = 'Training Loss') 20 | plt.plot(val_loss_plot[x_start:,0], val_loss_plot[x_start:,1], label = 'Validation Loss') 21 | plt.legend() 22 | plt.xlabel('Epoch') 23 | plt.ylabel('Loss') 24 | plt.show(block=True) 25 | 26 | -------------------------------------------------------------------------------- /pytorch/plumeConfig.yaml: -------------------------------------------------------------------------------- 1 | # Configuration file for buoyant simulation. 2 | # The program overwrites parameters from conf and mconf dicts found 3 | # in (where the trained model is) with parameters found here. 4 | # They just impact the physics of the problem but leave the architecture 5 | # untouched. 6 | # This table is saved to disk, in , at every iteration and 7 | # is loaded when restarting a sim. 8 | # To restart precise and set restartSim to true. 9 | 10 | # res : Resolution of domain 11 | resX: 128 12 | resY: 128 13 | # restartSim : set to true to restart form checkpoint from outputFolder. 14 | restartSim: false 15 | # realTimePlot : set to true to plot results during simulation. 16 | realTimePlot: true 17 | # saveVTK : set to true to save results as structured VTK files 18 | # use Paraview to load those files. 19 | saveVTK: false 20 | # outputFolder : folder for storing results, checkpoints and config files 21 | outputFolder: "/output/folder" 22 | # simMethod: choose convnet or jacobi 23 | simMethod: "convnet" 24 | # modelDir : location of trained model 25 | modelDir: "/path/to/folder/fluidnet_cxx/trained_models/ScaleNet_ShortTerm_LongTermLoss" 26 | # modelFilename : name of trained model 27 | modelFilename: "convModel" 28 | # dt : simualation timestep 29 | dt: 0.1 30 | # statIter : iterations frequency for output 31 | statIter: 100 32 | # maxIter : maximum number of simulation iterations 33 | maxIter: 20000 34 | # maccormackStrength : used in semi-lagrangian MacCormack advection 35 | # when LT div is activated. 0.6 is a good value. If ~1, can lead to 36 | # high frequency artifacts. 37 | maccormackStrength: 0.6 38 | # buoyancyScale : Buoyancy forces scale 39 | # gravityScale : Gravity forces scale 40 | # Note: Manta and FluidNet divide gravity forces into "gravity" and "buoyancy" 41 | # They represent the two terms arising from Boussinesq approximation 42 | # rho*g = rho_0*g + delta_rho*g 43 | # (1) (2) 44 | # rho_0 being the average density and delta_rho local difference of density 45 | # w.r.t average density. 46 | # Mantaflow calls (1) gravity and (2) buoyancy and allows for different g's 47 | buoyancyScale: 0.25 48 | gravityScale: 0 49 | # Introduces a correcting factor in the denisty equation 50 | # from "A splitting method for incompressible flows with variable 51 | # density based on a pressure Poisson equation" (Guermond, Salgado). 52 | # Not really tested... Recommendation is to leave it as false. 53 | correctScalar: false 54 | # viscosity : introduces a viscous term in moment equation. 55 | # Algortihm taken from the book "Fluid Simulation for Computer Graphics" by 56 | # Bridson 57 | viscosity: 0 58 | # operatingDensity : When applying buoyancy, buoyancyScale is multiplied 59 | # by (density(i,j) - operatingDensity) 60 | operatingDensity: 0.0 61 | # pTol : convergence criterion for Jacobi method 62 | pTol: 0.00 63 | # jacobiIter : maximum number of iterations for Jacobi method 64 | jacobiIter: 200 65 | # gravityVec : orientation of the gravity vector when buoyancy/gravity 66 | # are present. 67 | gravityVec: 68 | x: 0.0 69 | y: -1.0 70 | z: 0.0 71 | # injectionDensity : density at buoyant plume inlet 72 | injectionDensity: 0.1 73 | # sourceRadius : radius of inlet zone w.r.t the center of x-axis 74 | sourceRadius: 0.145 75 | # injectionVelocity : inlet velocity 76 | injectionVelocity: 2 77 | -------------------------------------------------------------------------------- /pytorch/print_output.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import sys 3 | import argparse 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | import torch.optim as optim 9 | from torch.autograd import Variable 10 | import torch.utils.data 11 | 12 | import glob 13 | from shutil import copyfile 14 | import importlib.util 15 | 16 | import lib 17 | import lib.fluid as fluid 18 | from config import defaultConf # Dictionnary with configuration and model params. 19 | 20 | # Use: python3 print_output.py 21 | # e.g: python3 print_output.py data/model_test convModel 22 | # 23 | # Utility to print data, output or target fields, as well as errors. 24 | # It loads the state and runs one epoch with a batch size = 1. It can be used while 25 | # training, to have some visual help. 26 | 27 | #**************************** Load command line arguments ********************* 28 | assert (len(sys.argv) == 3), 'Usage: python3 print_output.py ' 29 | assert (glob.os.path.exists(sys.argv[1])), 'Directory ' + str(sys.argv[1]) + ' does not exists' 30 | 31 | #********************************** Define Config ****************************** 32 | 33 | #TODO: allow to overwrite params from the command line by parsing. 34 | 35 | conf = defaultConf.copy() 36 | conf['modelDir'] = sys.argv[1] 37 | print(sys.argv[1]) 38 | conf['modelDirname'] = conf['modelDir'] + '/' + conf['modelFilename'] 39 | resume = False 40 | 41 | #*********************************** Select the GPU **************************** 42 | print('Active CUDA Device: GPU', torch.cuda.current_device()) 43 | 44 | path = conf['modelDir'] 45 | path_list = path.split(glob.os.sep) 46 | saved_model_name = glob.os.path.join(*path_list[:-1], path_list[-2] + '_saved.py') 47 | temp_model = glob.os.path.join('lib', path_list[-2] + '_saved_print.py') 48 | copyfile(saved_model_name, temp_model) 49 | 50 | assert glob.os.path.isfile(temp_model), temp_model + ' does not exits!' 51 | spec = importlib.util.spec_from_file_location('model_saved', temp_model) 52 | model_saved = importlib.util.module_from_spec(spec) 53 | spec.loader.exec_module(model_saved) 54 | 55 | try: 56 | te = lib.FluidNetDataset(conf, 'te', save_dt=4, resume=resume) # Test instance of custom Dataset 57 | 58 | conf, mconf = te.createConfDict() 59 | 60 | print('==> overwriting conf and file_mconf') 61 | cpath = glob.os.path.join(conf['modelDir'], conf['modelFilename'] + '_conf.pth') 62 | mcpath = glob.os.path.join(conf['modelDir'], conf['modelFilename'] + '_mconf.pth') 63 | assert glob.os.path.isfile(mcpath), cpath + ' does not exits!' 64 | assert glob.os.path.isfile(mcpath), mcpath + ' does not exits!' 65 | conf = torch.load(cpath) 66 | mconf = torch.load(mcpath) 67 | 68 | test_loader = torch.utils.data.DataLoader(te, batch_size=1, \ 69 | num_workers=0, shuffle=False, pin_memory=True) 70 | 71 | print('==> loading checkpoint') 72 | mpath = glob.os.path.join(conf['modelDir'], conf['modelFilename'] + '_lastEpoch_best.pth') 73 | assert glob.os.path.isfile(mpath), mpath + ' does not exits!' 74 | state = torch.load(mpath) 75 | 76 | print('Data loading: done') 77 | 78 | #********************************** Create the model *************************** 79 | 80 | print('') 81 | print('----- Model ------') 82 | 83 | # Create model and print layers and params 84 | 85 | net = model_saved.FluidNet(mconf, dropout=False) 86 | if torch.cuda.is_available(): 87 | net = net.cuda() 88 | #lib.summary(net, (3,1,128,128)) 89 | 90 | net.load_state_dict(state['state_dict']) 91 | 92 | #********************************* Run the model and print **************************** 93 | from itertools import count 94 | 95 | batch_print = [100, 1000, 4000, 5000, 8000, 10000, 15000, 20000] 96 | def val(): 97 | net.eval() 98 | 99 | #initialise loss scores 100 | total_loss = 0 101 | p_l2_total_loss = 0 102 | div_l2_total_loss = 0 103 | p_l1_total_loss = 0 104 | div_l1_total_loss = 0 105 | div_lt_total_loss = 0 106 | 107 | n_batches = 0 # Number processed 108 | 109 | # Loss types 110 | _pL2Loss = nn.MSELoss() 111 | _divL2Loss = nn.MSELoss() 112 | _divLTLoss = nn.MSELoss() 113 | _pL1Loss = nn.L1Loss() 114 | _divL1Loss = nn.L1Loss() 115 | 116 | # Loss lambdas (multiply the corresponding loss) 117 | pL2Lambda = mconf['pL2Lambda'] 118 | divL2Lambda = mconf['divL2Lambda'] 119 | pL1Lambda = mconf['pL1Lambda'] 120 | divL1Lambda = mconf['divL1Lambda'] 121 | divLTLambda = mconf['divLongTermLambda'] 122 | 123 | for batch_idx, (data, target) in zip(count(step=1), test_loader): 124 | with torch.no_grad(): 125 | if torch.cuda.is_available(): 126 | data, target = data.cuda(), target.cuda() 127 | if batch_idx in batch_print: 128 | out_p, out_U = net(data) 129 | flags = data[:,3].unsqueeze(1).contiguous() 130 | target_p = target[:,0].unsqueeze(1) 131 | out_div = fluid.velocityDivergence(\ 132 | out_U.contiguous(), \ 133 | flags) 134 | target_div = torch.zeros_like(out_div) 135 | 136 | # Measure loss and save it 137 | pL2Loss = pL2Lambda *_pL2Loss(out_p, target_p) 138 | divL2Loss = divL2Lambda *_divL2Loss(out_div, target_div) 139 | pL1Loss = pL1Lambda *_pL1Loss(out_p, target_p) 140 | divL1Loss = divL1Lambda *_divL1Loss(out_div, target_div) 141 | 142 | loss_size = pL2Loss + divL2Loss + pL1Loss + divL1Loss 143 | 144 | # Print statistics 145 | p_l2_total_loss += pL2Loss.data.item() 146 | div_l2_total_loss += divL2Loss.data.item() 147 | p_l1_total_loss += pL1Loss.data.item() 148 | div_l1_total_loss += divL1Loss.data.item() 149 | #if (divLTLambda > 0): 150 | # div_lt_total_loss += divLTLoss.data.item() 151 | total_loss += loss_size.data.item() 152 | 153 | 154 | # Measure loss and save it 155 | lib.plotField(out=[out_p, out_U, out_div], 156 | tar=target, 157 | flags=flags, 158 | loss=[total_loss, p_l2_total_loss, 159 | div_l2_total_loss, div_lt_total_loss, 160 | p_l1_total_loss, div_l1_total_loss], 161 | mconf=mconf, 162 | save=False, 163 | y_slice=8) 164 | 165 | #********************************* Run one epoch *************************************** 166 | 167 | print('Plotting results at epoch ' + str(state['epoch']) ) 168 | # Train on training set and test on validation set 169 | val() 170 | 171 | finally: 172 | # Delete model_saved.py 173 | print() 174 | print('Deleting' + temp_model) 175 | glob.os.remove(temp_model) 176 | -------------------------------------------------------------------------------- /pytorch/rayleighTaylorConfig.yaml: -------------------------------------------------------------------------------- 1 | # Configuration file for buoyant simulation. 2 | # The program overwrites parameters from conf and mconf dicts found 3 | # in (where the trained model is) with parameters found here. 4 | # They just impact the physics of the problem but leave the architecture 5 | # untouched. 6 | # This table is saved to disk, in , at every iteration and 7 | # is loaded when restarting a sim. 8 | # To restart precise and set restartSim to true. 9 | 10 | # res : Resolution of domain 11 | resX: 128 12 | resY: 512 13 | # restartSim : set to true to restart form checkpoint from outputFolder. 14 | restartSim: false 15 | # realTimePlot : set to true to plot results during simulation. 16 | realTimePlot: false 17 | # saveVTK : set to true to save results as structured VTK files 18 | # use Paraview to load those files. 19 | saveVTK: false 20 | # outputFolder : folder for storing results, checkpoints and config files 21 | outputFolder: "/path/to/output/folder" 22 | # simMethod: choose convnet or jacobi 23 | simMethod: "convnet" 24 | # modelDir : location of trained model 25 | modelDir: "/path/to/model/directory" 26 | # modelFilename : name of trained model 27 | modelFilename: "convModel" 28 | # dt : simualation timestep 29 | dt: 0.5 30 | # statIter : iterations frequency for output 31 | statIter: 10 32 | # maxIter : maximum number of simulation iterations 33 | maxIter: 20000 34 | # maccormackStrength : used in semi-lagrangian MacCormack advection 35 | # when LT div is activated. 0.6 is a good value. If ~1, can lead to 36 | # high frequency artifacts. 37 | maccormackStrength: 0.6 38 | # buoyancyScale : Buoyancy forces scale 39 | # gravityScale : Gravity forces scale 40 | # Note: Manta and FluidNet divide gravity forces into "gravity" and "buoyancy" 41 | # They represent the two terms arising from Boussinesq approximation 42 | # rho*g = rho_0*g + delta_rho*g 43 | # (1) (2) 44 | # rho_0 being the average density and delta_rho local difference of density 45 | # w.r.t average density. 46 | # Mantaflow calls (1) gravity and (2) buoyancy and allows for different g's 47 | buoyancyScale: 1. 48 | gravityScale: 0 49 | # Introduces a correcting factor in the denisty equation 50 | # from "A splitting method for incompressible flows with variable 51 | # density based on a pressure Poisson equation" (Guermond, Salgado). 52 | # Not really tested... Recommendation is to leave it as false. 53 | correctScalar: false 54 | # viscosity : introduces a viscous term in moment equation. 55 | # Algortihm taken from the book "Fluid Simulation for Computer Graphics" by 56 | # Bridson 57 | viscosity: 0 58 | # operatingDensity : When applying buoyancy, buoyancyScale is multiplied 59 | # by (density(i,j) - operatingDensity) 60 | operatingDensity: 0.0 61 | # pTol : convergence criterion for Jacobi method 62 | pTol: 0.0 63 | # jacobiIter : maximum number of iterations for Jacobi method 64 | jacobiIter: 200 65 | # gravityVec : orientation of the gravity vector when buoyancy/gravity 66 | # are present. 67 | gravityVec: 68 | x: 0.0 69 | y: 1.0 70 | z: 0.0 71 | # Rayleigh Taylor parameters 72 | rho2: 0.01 73 | rho1: -0.01 74 | perturbThickness: 100 75 | perturbAmplitude : 0.01 76 | height: 0.5 77 | -------------------------------------------------------------------------------- /pytorch/trainConfig.yaml: -------------------------------------------------------------------------------- 1 | # Configuration file with default parameters. 2 | # Some can be modified through the command line. See help function for training 3 | # script and README.md for more info. 4 | # This table is saved to disk (as pytorch objects) on every epoch 5 | # so that simulations can be paused and restarted. 6 | #========================================= 7 | # MODEL 8 | #========================================= 9 | 10 | #========================================= 11 | # DATA 12 | #========================================= 13 | # dataDir : Dataset location 14 | dataDir: "/path/to/data/dir" 15 | # dataset : Dataset name. Folder inside dataDir with training and testing scenes 16 | dataset: "output_current_model_sphere" 17 | # numWorkers : number of parallel workers for dataloader. Set to 0 to allow PyTorch 18 | # to automatically manage loading. 19 | numWorkers: 3 20 | # If true, dataset is preprocessed and programs exists. 21 | # Preprocessing is automatic if no previous preproc is detected on current dataset. 22 | preprocOriginalFluidNetDataOnly: false 23 | # shuffleTraining : Shuffles dataset 24 | shuffleTraining: true 25 | 26 | #========================================= 27 | # OUTPUT 28 | #========================================= 29 | # modelDir : Output folder for trained model and loss log. 30 | modelDir: "/path/to/output" 31 | # modelFilename : Trained model name 32 | modelFilename: "convModel" 33 | 34 | #========================================= 35 | # TRAINING MONITORING 36 | #========================================= 37 | 38 | # freqToFile : Epoch frequency for loss output to file/image saving. 39 | freqToFile: 5 40 | # printTraining : Debug options for training. 41 | # Prints or shows validation dataset and compares net 42 | # output to GT. 43 | # Options: save (save figures), show (shows in windows), none 44 | printTraining: "save" 45 | 46 | #========================================= 47 | # TRAINING PARAMETER 48 | #========================================= 49 | batchSize: 64 50 | # maxEpochs : Maximum number of epochs 51 | maxEpochs: 400 52 | # resume : resume training from checkpoint. 53 | resumeTraining: false 54 | modelParam: 55 | # model : options ('FluidNet', 'ScaleNet') 56 | # -FluidNet : uses the architecture found in lib/model.py (based on FluidNet) 57 | # -ScaleNet : uses a multiscale architecture found in lib/multi_scale_net.py 58 | model: "FluidNet" 59 | 60 | # inputChannels : Network inputs. At least one of them must be set to true! 61 | inputChannels: 62 | div: true 63 | pDiv: false 64 | UDiv: false 65 | # lr : learning rate. If using scientific notation, necessary to precise type 66 | # for yaml->python cast. 67 | lr: !!python/float 5e-5 68 | # fooLambda : Weighting for each loss. Set to 0 to disable loss. 69 | # MSE of pressure 70 | pL2Lambda: 0 71 | # MSE of divergence (Ground truth is zero divergence) 72 | divL2Lambda: 1 73 | # Absolute difference of pressure 74 | pL1Lambda: 0 75 | # Absolute difference of divergence 76 | divL1Lambda: 0 77 | # MSE of long term divergence 78 | # If > 0, implements the Long Term divergence concept from FluidNet 79 | divLongTermLambda: 1 80 | # longTermDivNumSteps : We want to measure what the divergence is after 81 | # a set number of steps for each training and test sample. Set table 82 | # to nil to disable, (or set longTermDivLambda to 0). 83 | longTermDivNumSteps: 84 | - 4 85 | - 16 86 | # longTermDivProbability is the probability that longTermDivNumSteps[0] 87 | # will be taken, otherwise longTermDivNumSteps[1] will be taken with 88 | # probability of 1 - longTermDivProbability. 89 | longTermDivProbability: 0.9 90 | # normalizeInput : if true, normalizes input by max(std(chan), threshold) 91 | normalizeInput: true 92 | # normalizeInputChan : which channel to calculate std 93 | normalizeInputChan: "UDiv" 94 | # normalizeInputThreshold : don't normalize input noise 95 | normalizeInputThreshold: 0.00001 96 | 97 | #========================================= 98 | # PHYSICAL PARAMETERS 99 | #========================================= 100 | # Time step: default simulation timestep. 101 | dt: 0.1 102 | 103 | # ONLY APPLIED IF LONG TERM DIV IS ACTIVATED 104 | # ---------------------------------- 105 | # buoyancyScale : Buoyancy forces scale 106 | # gravityScale : Gravity forces scale 107 | # Note: Manta and FluidNet divide gravity forces into "gravity" and "buoyancy" 108 | # They represent the two terms arising from Boussinesq approximation 109 | # rho*g = rho_0*g + delta_rho*g 110 | # (1) (2) 111 | # rho_0 being the average density and delta_rho local difference of density 112 | # w.r.t average density. 113 | # Mantaflow calls (1) gravity and (2) buoyancy and allows for different g's 114 | buoyancyScale: 0 115 | gravityScale: 0 116 | # Gravity vector: Direction of gravity Vector 117 | gravityVec: 118 | x: 0 119 | y: 0 120 | z: 0 121 | # training buoyancy scale : This is the buoyancy to use when adding buoyancy 122 | # to the long term training. It will be applied in a random cardinal direction. 123 | trainBuoyancyScale: 2.0 124 | # training buoyancy probability : This is the probability to add buoyancy when 125 | # long term training. 126 | trainBuoyancyProb: 0.3 127 | # training gravity scale : This is the gravity to use when adding gravity 128 | # to the long term training. It will be applied in a random cardinal direction. 129 | trainGravityScale: 0.0 130 | # training gravity probability : This is the probability to add buoyancy when 131 | # long term training. 132 | trainGravityProb: 0.0 133 | # ------------------------------------ 134 | # Introduces a correcting factor in the denisty equation 135 | # from "A splitting method for incompressible flows with variable 136 | # density based on a pressure Poisson equation" (Guermond, Salgado). 137 | # Not really tested... Recommendation is to leave it as false. 138 | correctScalar: false 139 | # operatingDensity : When applying buoyancy, buoyancyScale is multiplied 140 | # by (density(i,j) - operatingDensity) 141 | operatingDensity: 0.0 142 | # viscosity : introduces a viscous term in moment equation. 143 | # Algortihm taken from the book "Fluid Simulation for Computer Graphics" by 144 | # Bridson 145 | viscosity: 0 146 | # timeScaleSigma : Amplitude of time scale perturb during training. 147 | timeScaleSigma: 1 148 | # maccormackStrength : used in semi-lagrangian MacCormack advection 149 | # when LT div is activated. 0.6 is a good value. If ~1, can lead to 150 | # high frequency artifacts. 151 | maccormackStrength: 0.6 152 | # sampleOutsideFluid : if true, allows particles in advection to 'land' inside 153 | # obstacles. In general, we don't want that, so leave it as false to avoid 154 | # possible artifacts. 155 | sampleOutsideFluid: false 156 | 157 | -------------------------------------------------------------------------------- /solver_cpp/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | CMAKE_MINIMUM_REQUIRED(VERSION 3.0) 2 | PROJECT("fluidnet_cxx" LANGUAGES CXX) 3 | 4 | INCLUDE(CMakeDependentOption) 5 | OPTION(FLUID_TEST "Build Fluid test binaries" OFF) 6 | 7 | SET(CMAKE_BUILD_TYPE Release) 8 | 9 | # Include manually ATen 10 | SET(ATEN_INCLUDE_DIR "/usr/local/lib/python3.5/dist-packages/torch/lib/include") 11 | LINK_DIRECTORIES("/usr/local/lib/python3.5/dist-packages/torch/lib/") 12 | INCLUDE_DIRECTORIES(${ATEN_INCLUDE_DIR}) 13 | 14 | SET(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -ggdb -std=c++11") 15 | 16 | FIND_PACKAGE(OpenCV REQUIRED) 17 | INCLUDE_DIRECTORIES(${OpenCV_INCLUDE_DIRS}) 18 | 19 | ADD_SUBDIRECTORY(src) 20 | IF (FLUID_TEST) 21 | ADD_SUBDIRECTORY(test) 22 | ELSE () 23 | ADD_SUBDIRECTORY(simulate) 24 | ENDIF() 25 | 26 | 27 | 28 | -------------------------------------------------------------------------------- /solver_cpp/simulate/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | SET(SIM_SRC simulate.cpp) 2 | 3 | ADD_EXECUTABLE(fluidnet_sim ${SIM_SRC}) 4 | 5 | TARGET_INCLUDE_DIRECTORIES(fluidnet_sim PUBLIC 6 | $ 7 | ) 8 | LINK_DIRECTORIES(fluidnet_sim $) 9 | 10 | TARGET_LINK_LIBRARIES(fluidnet_sim FluidNet ATen) 11 | TARGET_LINK_LIBRARIES(fluidnet_sim ${OpenCV_LIBS} ) 12 | -------------------------------------------------------------------------------- /solver_cpp/src/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | IF (FLUID_TEST) 2 | SET(TFLUIDS 3 | fluidnet_implementation/init.cpp 4 | fluidnet_implementation/advect_type.cpp 5 | ) 6 | ELSE () 7 | SET(TFLUIDS ) 8 | ENDIF() 9 | 10 | SET(GRID 11 | grid/bool_conversion.cpp 12 | grid/grid.cpp 13 | ) 14 | SET(ADVECTION 15 | advection/advection.cpp 16 | advection/advect_type.cpp 17 | advection/calc_line_trace.cpp 18 | ) 19 | SET(BCS 20 | boundaryCondition/bcs.cpp 21 | ) 22 | SET(STERMS 23 | sourceTerms/source_term.cpp 24 | ) 25 | SET(PROJECTION 26 | projection/div.cpp 27 | projection/solve_linear_sys.cpp 28 | projection/update_vel.cpp 29 | ) 30 | ADD_LIBRARY(FluidNet SHARED ${TFLUIDS} ${GRID} ${ADVECTION} ${ADVECTION_OLD} ${BCS} ${STERMS} ${PROJECTION}) 31 | TARGET_INCLUDE_DIRECTORIES(FluidNet 32 | PUBLIC 33 | $ 41 | ) 42 | 43 | -------------------------------------------------------------------------------- /solver_cpp/src/advection/.gitignore: -------------------------------------------------------------------------------- 1 | debug 2 | old 3 | -------------------------------------------------------------------------------- /solver_cpp/src/advection/advect_type.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "ATen/ATen.h" 3 | #include "advect_type.h" 4 | 5 | AdvectMethod StringToAdvectMethod(const std::string& str) { 6 | if (str == "eulerFluidNet") { 7 | return ADVECT_EULER_FLUIDNET; 8 | } else if (str == "maccormackFluidNet") { 9 | return ADVECT_MACCORMACK_FLUIDNET; 10 | } else { 11 | std::stringstream ss; 12 | ss << "advection method (" << str << ") not supported (options " 13 | << "are: eulerFluidNet, maccormackFluidNet)"; 14 | AT_ERROR("Error: Advection method not supported"); 15 | } 16 | } 17 | 18 | -------------------------------------------------------------------------------- /solver_cpp/src/advection/advect_type.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | typedef enum { 6 | ADVECT_EULER_FLUIDNET = 0, 7 | ADVECT_MACCORMACK_FLUIDNET = 1, 8 | } AdvectMethod; 9 | 10 | AdvectMethod StringToAdvectMethod(const std::string& str); 11 | -------------------------------------------------------------------------------- /solver_cpp/src/advection/advection.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #include "ATen/ATen.h" 6 | #include "grid/grid.h" 7 | #include "grid/cell_type.h" 8 | #include "advect_type.h" 9 | #include "calc_line_trace.h" 10 | 11 | namespace fluid { 12 | 13 | typedef at::Tensor T; 14 | 15 | T SemiLagrangeEulerFluidNet 16 | ( 17 | T& flags, T& vel, T& src, T& maskBorder, 18 | float dt, float order_space, 19 | T& i, T& j, T& k, 20 | const bool line_trace, 21 | const bool sample_outside_fluid 22 | ); 23 | 24 | T SemiLagrangeEulerFluidNetSavePos 25 | ( 26 | T& flags, T& vel, T& src, T& maskBorder, 27 | float dt, float order_space, 28 | T& i, T& j, T& k, 29 | const bool line_trace, 30 | const bool sample_outside_fluid, 31 | T& pos 32 | ); 33 | 34 | T MacCormackCorrect 35 | ( 36 | T& flags, const T& old, 37 | const T& fwd, const T& bwd, 38 | const float strength, 39 | bool is_levelset 40 | ); 41 | 42 | T getClampBounds 43 | ( 44 | const T& src, const T& pos, const T& flags, 45 | const bool sample_outside, 46 | T& clamp_min, T& clamp_max 47 | ); 48 | 49 | T MacCormackClampFluidNet( 50 | T& flags, T& vel, 51 | const T& dst, const T& src, 52 | const T& fwd, const T& fwd_pos, 53 | const T& bwd_pos, const bool sample_outside 54 | ); 55 | 56 | // Advect scalar field 'p' by the input vel field 'u'. 57 | // 58 | // @input dt - timestep (seconds). 59 | // @input s - input scalar field to advect 60 | // @input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 61 | // @input flags - input occupancy grid 62 | // @input sDst - Returned scalar field. 63 | // @input method - OPTIONAL - "eulerFluidNet", "maccormackFluidNet" 64 | // @param boundaryWidth - OPTIONAL - boundary width. (default 1) 65 | // @param sampleOutsideFluid - OPTIONAL - For density advection we do not want 66 | // to advect values inside non-fluid cells and so this should be set to false. 67 | // For other quantities (like temperature), this should be true. 68 | // @param maccormackStrength - OPTIONAL - (default 0.75) A strength parameter 69 | // will make the advection eularian (with values interpolating in between). A 70 | // value of 1 (which implements the update from An Unconditionally Stable 71 | // MaCormack Method) tends to add too much high-frequency detail 72 | void advectScalar 73 | ( 74 | float dt, T& src, T& U, T& flags, T& s_dst, 75 | const std::string method_str = "maccormackFluidNet", 76 | int bnd = 1, 77 | const bool sample_outside_fluid = false, 78 | const float maccormack_strength = 0.75 79 | ); 80 | 81 | T SemiLagrangeEulerFluidNetMAC 82 | ( 83 | T& flags, T& vel, T& src, T& mask, 84 | float dt, float order_space, 85 | const bool line_trace, 86 | T& i, T& j, T& k 87 | ); 88 | 89 | T MacCormackCorrectMAC 90 | ( 91 | T& flags, const T& old, 92 | const T& fwd, const T& bwd, 93 | const float strength, 94 | T& i, T& j, T& k 95 | ); 96 | 97 | T doClampComponentMAC 98 | ( 99 | int chan, 100 | const T& flags, const T& dst, 101 | const T& orig, const T& fwd, 102 | const T& pos, const T& vel 103 | ); 104 | 105 | T MacCormackClampMAC 106 | ( 107 | const T& flags, const T& vel, const T& dval, 108 | const T& orig, const T& fwd, const T& mask, 109 | float dt, 110 | const T& i, const T& j, const T& k 111 | ); 112 | 113 | // Advect velocity field 'u' by itself and store in uDst. 114 | // 115 | // @input dt - timestep (seconds). 116 | // @input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 117 | // @input flags - input occupancy grid 118 | // @input UDst - Returned velocity field. 119 | // @input method - OPTIONAL - "eulerFluidNet", "maccormackFluidNet" (default) 120 | // @input boundaryWidth - OPTIONAL - boundary width. (default 1) 121 | // @input maccormackStrength - OPTIONAL - (default 0.75) A strength parameter 122 | // will make the advection more 1st order (with values interpolating in 123 | // between). A value of 1 (which implements the update from "An Unconditionally 124 | // Stable MaCormack Method") tends to add too much high-frequency detail. 125 | void advectVel 126 | ( 127 | float dt, T& U, T& flags, T& U_dst, 128 | const std::string method_str = "maccormackFluidNet", 129 | int bnd = 1, 130 | const float maccormack_strength = 0.75 131 | ); 132 | 133 | } // namespace fluid 134 | -------------------------------------------------------------------------------- /solver_cpp/src/advection/calc_line_trace.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "ATen/ATen.h" 4 | #include "grid/cell_type.h" 5 | 6 | // See the .cpp file for detailed comments. 7 | namespace fluid { 8 | 9 | typedef at::Tensor T; 10 | 11 | void getPixelCenter(const T& pos, T& ix); 12 | 13 | T isOutOfDomain(const T& pos, const T& flags); 14 | 15 | T isBlockedCell(const T& pos, const T& flags); 16 | 17 | void clampToDomain(T& pos, const T& flags); 18 | 19 | typedef enum Quadrants { 20 | RIGHT = 0, 21 | LEFT = 1, 22 | MIDDLE = 2 23 | } Quadrants; 24 | 25 | T HitBoundingBox(const T& minB, const T& maxB, 26 | const T& origin, const T& dir, 27 | const T& mask, T& coord); 28 | 29 | T calcRayBoxIntersection(const T& pos, const T& dt, const T& ctr, 30 | const float hit_margin, const T& mask, T& ipos); 31 | 32 | T calcRayBorderIntersection(const T& pos, const T& next_pos, 33 | const T& flags, const float hit_margin, 34 | const T& mOutDom, 35 | T& ipos); 36 | 37 | void calcLineTrace(const T& pos, const T& delta, const T& flags, 38 | T& new_pos, const bool do_line_trace); 39 | 40 | } // namespace fluid 41 | -------------------------------------------------------------------------------- /solver_cpp/src/boundaryCondition/.gitignore: -------------------------------------------------------------------------------- 1 | old 2 | -------------------------------------------------------------------------------- /solver_cpp/src/boundaryCondition/bcs.cpp: -------------------------------------------------------------------------------- 1 | #include "bcs.h" 2 | 3 | namespace fluid { 4 | 5 | typedef at::Tensor T; 6 | 7 | // ***************************************************************************** 8 | // setWallBcsForward 9 | // ***************************************************************************** 10 | 11 | // Enforce boundary conditions on velocity MAC Grid (i.e. set slip components). 12 | // 13 | // @input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 14 | // @input flags - input occupancy grid 15 | 16 | void setWallBcsForward 17 | ( 18 | T& U, T& flags 19 | ) { 20 | // Check arguments. 21 | AT_ASSERT(U.dim() == 5 && flags.dim() == 5, "Dimension mismatch"); 22 | AT_ASSERT(flags.size(1) == 1, "flags is not scalar"); 23 | int bsz = flags.size(0); 24 | int d = flags.size(2); 25 | int h = flags.size(3); 26 | int w = flags.size(4); 27 | 28 | bool is3D = (U.size(1) == 3); 29 | if (!is3D) { 30 | AT_ASSERT(d == 1, "2D velocity field but zdepth > 1"); 31 | AT_ASSERT(U.size(1) == 2, "2D velocity field must have only 2 channels"); 32 | } 33 | AT_ASSERT((U.size(0) == bsz && U.size(2) == d && 34 | U.size(3) == h && U.size(4) == w), "Size mismatch"); 35 | 36 | AT_ASSERT(U.is_contiguous() && flags.is_contiguous(), 37 | "Input is not contiguous"); 38 | 39 | T i = infer_type(flags).arange(0, w).view({1,w}).expand({bsz, d, h, w}).toType(at::kLong); 40 | T j = infer_type(i).arange(0, h).view({1,h,1}).expand({bsz, d, h, w}); 41 | T k = zeros_like(i); 42 | if (is3D) { 43 | k = infer_type(i).arange(0, d).view({1,d,1,1}).expand({bsz, d, h, w}); 44 | } 45 | T zero = zeros_like(i); 46 | T zeroBy = zero.toType(at::kByte); 47 | 48 | T idx_b = infer_type(i).arange(0, bsz).view({bsz,1,1,1}); 49 | idx_b = idx_b.expand({bsz,d,h,w}); 50 | 51 | T mCont = ones_like(zeroBy); 52 | 53 | T cur_fluid = flags.eq(TypeFluid).squeeze(1); 54 | T cur_obs = flags.eq(TypeObstacle).squeeze(1); 55 | T mNotFluidNotObs = cur_fluid.ne(1).__and__(cur_obs.ne(1)); 56 | mCont.masked_fill_(mNotFluidNotObs, 0); 57 | 58 | T i_l = zero.where( (i <=0), i - 1); 59 | T obst100 = zeroBy.where( i <= 0, (flags.index({idx_b, zero, k, j, i_l}).eq(TypeObstacle))).__and__(mCont); 60 | U.select(1,0).masked_fill_(obst100, 0); 61 | 62 | T obs_fluid100 = zeroBy.where( i <= 0, (flags.index({idx_b, zero, k, j, i_l}).eq(TypeFluid))). 63 | __and__(cur_obs).__and__(mCont); 64 | U.select(1,0).masked_fill_(obs_fluid100, 0); 65 | 66 | T j_l = zero.where( (j <= 0), j - 1); 67 | T obst010 = zeroBy.where( j <= 0, (flags.index({idx_b, zero, k, j_l, i}).eq(TypeObstacle))).__and__(mCont); 68 | U.select(1,1).masked_fill_(obst010, 0); 69 | 70 | T obs_fluid010 = zeroBy.where( j <= 0, (flags.index({idx_b, zero, k, j_l, i}).eq(TypeFluid))). 71 | __and__(cur_obs).__and__(mCont); 72 | U.select(1,1).masked_fill_(obs_fluid010, 0); 73 | 74 | if (is3D) { 75 | T k_l = zero.where( (k <= 0), k - 1); 76 | 77 | T obst001 = zeroBy.where( k <= 0, (flags.index({idx_b, zero, k_l, j, i}).eq(TypeObstacle))).__and__(mCont); 78 | U.select(1,2).masked_fill_(obst001, 0); 79 | 80 | T obs_fluid001 = zeroBy.where( k <= 0, (flags.index({idx_b, zero, k_l, j, i}).eq(TypeFluid))). 81 | __and__(cur_obs).__and__(mCont); 82 | U.select(1,2).masked_fill_(obs_fluid001, 0); 83 | } 84 | 85 | // TODO: implement TypeStick BCs 86 | } 87 | 88 | } // namespace fluid 89 | -------------------------------------------------------------------------------- /solver_cpp/src/boundaryCondition/bcs.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "ATen/ATen.h" 4 | #include "grid/grid.h" 5 | #include "grid/cell_type.h" 6 | 7 | namespace fluid { 8 | 9 | typedef at::Tensor T; 10 | 11 | // ***************************************************************************** 12 | // setWallBcsForward 13 | // ***************************************************************************** 14 | 15 | // Enforce boundary conditions on velocity MAC Grid (i.e. set slip components). 16 | // 17 | // @input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 18 | // @input flags - input occupancy grid 19 | 20 | void setWallBcsForward 21 | ( 22 | T& U, 23 | T& flags 24 | ); 25 | 26 | } // namespace fluid 27 | -------------------------------------------------------------------------------- /solver_cpp/src/fluid.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "grid/grid.h" 4 | #include "grid/cell_type.h" 5 | #include "advection/advection.h" 6 | #include "boundaryCondition/bcs.h" 7 | #include "sourceTerms/source_term.h" 8 | #include "projection/div.h" 9 | #include "projection/solve_linear_sys.h" 10 | #include "projection/update_vel.h" 11 | 12 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | CMAKE_MINIMUM_REQUIRED(VERSION 3.0) 2 | PROJECT("advection" LANGUAGES CXX) 3 | 4 | SET(CMAKE_BUILD_TYPE Release) 5 | 6 | # Include manually ATen 7 | SET(ATEN_INCLUDE_DIR "/usr/local/lib/python3.5/dist-packages/torch/lib/include") 8 | LINK_DIRECTORIES("/usr/local/lib/python3.5/dist-packages/torch/lib/") 9 | INCLUDE_DIRECTORIES(${ATEN_INCLUDE_DIR}) 10 | 11 | FIND_PACKAGE(OpenMP) 12 | IF (OPENMP_FOUND) 13 | SET (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS} -Wall -ggdb -std=c++11") 14 | SET (CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} ${OpenMP_EXE_LINKER_FLAGS}") 15 | ENDIF() 16 | 17 | 18 | SET(SIM_SRC init.cpp) 19 | 20 | ADD_EXECUTABLE(fluidnet_sim ${SIM_SRC}) 21 | 22 | TARGET_INCLUDE_DIRECTORIES(fluidnet_sim PUBLIC 23 | $ 24 | ) 25 | LINK_DIRECTORIES(fluidnet_sim $) 26 | 27 | TARGET_LINK_LIBRARIES(fluidnet_sim ATen) 28 | 29 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/advect_type.cpp: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include 16 | #include "advect_type.h" 17 | 18 | namespace tfluids { 19 | AdvectMethod StringToAdvectMethod(const std::string& str) { 20 | if (str == "euler") { 21 | return ADVECT_EULER_MANTA; 22 | } else if (str == "maccormack") { 23 | return ADVECT_MACCORMACK_MANTA; 24 | } else if (str == "rk2Ours") { 25 | return ADVECT_RK2_OURS; 26 | } else if (str == "eulerOurs") { 27 | return ADVECT_EULER_OURS; 28 | } else if (str == "rk3Ours") { 29 | return ADVECT_RK3_OURS; 30 | } else if (str == "maccormackFluidNet") { 31 | return ADVECT_MACCORMACK_OURS; 32 | } else { 33 | std::stringstream ss; 34 | ss << "advection method (" << str << ") not supported (options " 35 | << "are: euler, maccormack, rk2Ours, rk3Ours, eulerOurs)"; 36 | AT_ERROR(ss.str().c_str()); 37 | } 38 | } 39 | } // namespace tfluids 40 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/advect_type.h: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #pragma once 16 | 17 | #include 18 | #include "ATen/ATen.h" 19 | 20 | namespace tfluids{ 21 | typedef enum { 22 | ADVECT_EULER_MANTA = 0, // Port of Manta's implementation. 23 | ADVECT_MACCORMACK_MANTA = 1, // Port of Manta's implementation. 24 | ADVECT_EULER_OURS = 2, 25 | ADVECT_RK2_OURS = 3, 26 | ADVECT_RK3_OURS = 4, 27 | ADVECT_MACCORMACK_OURS = 5, 28 | } AdvectMethod; 29 | 30 | AdvectMethod StringToAdvectMethod(const std::string& str); 31 | 32 | } // namespace tfluids 33 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/cell_type.h: -------------------------------------------------------------------------------- 1 | // The methods in this file borrow heavily from - or or a direct port of - parts 2 | // of the Mantaflow library. Since the Mantaflow library is GNU GPL V3, we are 3 | // releasing this code under GNU as well as a third_party add on. See 4 | // FluidNet/torch/tfluids/third_party/README for more information. 5 | 6 | /****************************************************************************** 7 | * 8 | * MantaFlow fluid solver framework 9 | * Copyright 2011 Tobias Pfaff, Nils Thuerey 10 | * 11 | * This program is free software, distributed under the terms of the 12 | * GNU General Public License (GPL) 13 | * http://www.gnu.org/licenses 14 | * 15 | ******************************************************************************/ 16 | 17 | #pragma once 18 | 19 | // These are the same enum values used in Manta. We can't include grid.h 20 | // from Manta without pulling in the entire library, so we'll just redefine 21 | // them here. 22 | enum CellType { 23 | TypeNone = 0, 24 | TypeFluid = 1, 25 | TypeObstacle = 2, 26 | TypeEmpty = 4, 27 | TypeInflow = 8, 28 | TypeOutflow = 16, 29 | TypeOpen = 32, 30 | TypeStick = 128, 31 | TypeReserved = 256, 32 | TypeZeroPressure = (1<<15) 33 | }; 34 | 35 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/init.cpp: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include 16 | #include 17 | #include "ATen/ATen.h" 18 | 19 | #include "stack_trace.cpp" 20 | #include "cell_type.h" 21 | 22 | // This type is common to both float and double implementations and so has 23 | // to be defined outside tfluids.cc. 24 | #include "int3.h" 25 | #include "advect_type.h" 26 | 27 | // Some common functions 28 | inline int32_t clamp(const int32_t x, const int32_t low, const int32_t high) { 29 | return std::max(std::min(x, high), low); 30 | } 31 | 32 | // Expand the CPU types (float and double). This actually instantiates the 33 | // functions. Note: the order here is important. 34 | 35 | #define tfluids_(NAME) tfluids_ ## Real ## NAME 36 | 37 | #define real float 38 | #define accreal double 39 | #define Real Float 40 | #define TH_REAL_IS_FLOAT 41 | #include "vec3.cpp" 42 | #include "grid.cpp" 43 | #include "tfluids.cpp" 44 | #undef accreal 45 | #undef real 46 | #undef Real 47 | #undef THInf 48 | #undef TH_REAL_IS_FLOAT 49 | 50 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/int3.h: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #pragma once 16 | 17 | #include 18 | #include 19 | 20 | struct Int3 { 21 | int32_t x; 22 | int32_t y; 23 | int32_t z; 24 | 25 | Int3() : x(0), y(0), z(0) { } 26 | Int3(int32_t _x, int32_t _y, int32_t _z) : 27 | x(_x), y(_y), z(_z) { } 28 | 29 | Int3& operator=(const Int3& other) { 30 | if (this != &other) { 31 | this->x = other.x; 32 | this->y = other.y; 33 | this->z = other.z; 34 | } 35 | return *this; 36 | } 37 | 38 | Int3& operator+=(const Int3& rhs) { 39 | this->x += rhs.x; 40 | this->y += rhs.y; 41 | this->z += rhs.z; 42 | return *this; 43 | } 44 | 45 | const Int3 operator+(const Int3& rhs) const { 46 | Int3 ret = *this; 47 | ret += rhs; 48 | return ret; 49 | } 50 | 51 | Int3& operator-=(const Int3& rhs) { 52 | this->x -= rhs.x; 53 | this->y -= rhs.y; 54 | this->z -= rhs.z; 55 | return *this; 56 | } 57 | 58 | const Int3 operator-(const Int3& rhs) const { 59 | Int3 ret = *this; 60 | ret -= rhs; 61 | return ret; 62 | } 63 | 64 | const Int3 operator+(const int32_t rhs) const { 65 | Int3 ret = *this; 66 | ret.x += rhs; 67 | ret.y += rhs; 68 | ret.z += rhs; 69 | return ret; 70 | } 71 | 72 | const Int3 operator-(const int32_t rhs) const { 73 | Int3 ret = *this; 74 | ret.x -= rhs; 75 | ret.y -= rhs; 76 | ret.z -= rhs; 77 | return ret; 78 | } 79 | 80 | const Int3 operator*(const int32_t rhs) const { 81 | Int3 ret = *this; 82 | ret.x *= rhs; 83 | ret.y *= rhs; 84 | ret.z *= rhs; 85 | return ret; 86 | } 87 | }; 88 | 89 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/quadrants.h: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #ifndef QUADRANTS_H 16 | #define QUADRANTS_H 17 | 18 | typedef enum Quadrants { 19 | RIGHT = 0, 20 | LEFT = 1, 21 | MIDDLE = 2 22 | } Quadrants; 23 | 24 | #endif 25 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/stack_trace.cpp: -------------------------------------------------------------------------------- 1 | // In our grid classes it would be nice to be able to see the stacktrace if 2 | // we encounter an out of bounds access. 3 | 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | static std::mutex global_mutex; 14 | 15 | void PrintStacktrace() { 16 | std::lock_guard lock(global_mutex); 17 | const int max_num_frames = 64; 18 | void *frames[max_num_frames]; 19 | size_t num_frames = backtrace(frames, max_num_frames); 20 | 21 | // Print raw stack track, note: likely all symbol names will be mangled 22 | // but I can't figure out how to prevent this on nvcc + openmp (i.e. I can 23 | // get the first frame (this function) using abi::__cxa_demangle, but every 24 | // frame after that is hidden. 25 | backtrace_symbols_fd(frames, num_frames, STDERR_FILENO); 26 | } 27 | 28 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/vec3.cpp: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include 16 | #include 17 | 18 | #include "vec3.h" 19 | 20 | static Int3 toInt3(const tfluids_(vec3)& val) { 21 | return Int3(static_cast(val.x), 22 | static_cast(val.y), 23 | static_cast(val.z)); 24 | }; 25 | -------------------------------------------------------------------------------- /solver_cpp/src/fluidnet_implementation/vec3.h: -------------------------------------------------------------------------------- 1 | // Copyright 2016 Google Inc, NYU. 2 | // 3 | // Licensed under the Apache License, Version 2.0 (the "License"); 4 | // you may not use this file except in compliance with the License. 5 | // You may obtain a copy of the License at 6 | // 7 | // http://www.apache.org/licenses/LICENSE-2.0 8 | // 9 | // Unless required by applicable law or agreed to in writing, software 10 | // distributed under the License is distributed on an "AS IS" BASIS, 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | // See the License for the specific language governing permissions and 13 | // limitations under the License. 14 | 15 | #include 16 | #include 17 | 18 | struct tfluids_(vec3) { 19 | #ifdef TH_REAL_IS_FLOAT 20 | constexpr static const real kEpsilon = 1e-6f; 21 | #endif 22 | #ifdef TH_REAL_IS_DOUBLE 23 | constexpr static const real kEpsilon = 1e-10; 24 | #endif 25 | 26 | real x; 27 | real y; 28 | real z; 29 | 30 | tfluids_(vec3)() : x(0), y(0), z(0) { } 31 | tfluids_(vec3)(real _x, real _y, real _z) : x(_x), y(_y), z(_z) { } 32 | 33 | tfluids_(vec3)& operator=(const tfluids_(vec3)& other) { // Copy assignment. 34 | if (this != &other) { 35 | this->x = other.x; 36 | this->y = other.y; 37 | this->z = other.z; 38 | } 39 | return *this; 40 | } 41 | 42 | tfluids_(vec3)& operator+=(const tfluids_(vec3)& rhs) { // accum vec 43 | this->x += rhs.x; 44 | this->y += rhs.y; 45 | this->z += rhs.z; 46 | return *this; 47 | } 48 | 49 | const tfluids_(vec3) operator+(const tfluids_(vec3)& rhs) const { // add vec 50 | tfluids_(vec3) ret = *this; 51 | ret += rhs; 52 | return ret; 53 | } 54 | 55 | tfluids_(vec3)& operator-=(const tfluids_(vec3)& rhs) { // neg accum vec 56 | this->x -= rhs.x; 57 | this->y -= rhs.y; 58 | this->z -= rhs.z; 59 | return *this; 60 | } 61 | 62 | const tfluids_(vec3) operator-(const tfluids_(vec3)& rhs) const { // sub vec 63 | tfluids_(vec3) ret = *this; 64 | ret -= rhs; 65 | return ret; 66 | } 67 | 68 | const tfluids_(vec3) operator+(const real rhs) const { // add scalar 69 | tfluids_(vec3) ret = *this; 70 | ret.x += rhs; 71 | ret.y += rhs; 72 | ret.z += rhs; 73 | return ret; 74 | } 75 | 76 | const tfluids_(vec3) operator-(const real rhs) const { // sub scalar 77 | tfluids_(vec3) ret = *this; 78 | ret.x -= rhs; 79 | ret.y -= rhs; 80 | ret.z -= rhs; 81 | return ret; 82 | } 83 | 84 | const tfluids_(vec3) operator*(const real rhs) const { // mult scalar 85 | tfluids_(vec3) ret = *this; 86 | ret.x *= rhs; 87 | ret.y *= rhs; 88 | ret.z *= rhs; 89 | return ret; 90 | } 91 | 92 | const tfluids_(vec3) operator/(const real rhs) const { // mult scalar 93 | tfluids_(vec3) ret = *this; 94 | ret.x /= rhs; 95 | ret.y /= rhs; 96 | ret.z /= rhs; 97 | return ret; 98 | } 99 | 100 | inline real& operator()(int32_t i) { 101 | switch (i) { 102 | case 0: 103 | return this->x; 104 | case 1: 105 | return this->y; 106 | case 2: 107 | return this->z; 108 | default: 109 | AT_ERROR("vec3 out of bounds."); 110 | exit(-1); 111 | break; 112 | } 113 | } 114 | 115 | inline real operator()(int32_t i) const { 116 | return (*this)(i); 117 | } 118 | 119 | inline real norm() const { 120 | const real length_sq = 121 | this->x * this->x + this->y * this->y + this->z * this->z; 122 | if (length_sq > static_cast(kEpsilon)) { 123 | return std::sqrt(length_sq); 124 | } else { 125 | return static_cast(0); 126 | } 127 | } 128 | 129 | inline void normalize() { 130 | const real norm = this->norm(); 131 | if (norm > static_cast(kEpsilon)) { 132 | this->x /= norm; 133 | this->y /= norm; 134 | this->z /= norm; 135 | } else { 136 | this->x = 0; 137 | this->y = 0; 138 | this->z = 0; 139 | } 140 | } 141 | 142 | static tfluids_(vec3) cross(const tfluids_(vec3)& a, 143 | const tfluids_(vec3)& b) { 144 | tfluids_(vec3) ret; 145 | ret.x = (a.y * b.z) - (a.z * b.y); 146 | ret.y = (a.z * b.x) - (a.x * b.z); 147 | ret.z = (a.x * b.y) - (a.y * b.x); 148 | return ret; 149 | } 150 | }; 151 | 152 | static Int3 toInt3(const tfluids_(vec3)& val); 153 | -------------------------------------------------------------------------------- /solver_cpp/src/grid/.gitignore: -------------------------------------------------------------------------------- 1 | old 2 | -------------------------------------------------------------------------------- /solver_cpp/src/grid/bool_conversion.cpp: -------------------------------------------------------------------------------- 1 | #include "bool_conversion.h" 2 | 3 | namespace fluid { 4 | 5 | bool toBool(const at::Tensor & self) { 6 | return self.equal(self.type().ones({})); 7 | } 8 | 9 | } // namespace fluid 10 | -------------------------------------------------------------------------------- /solver_cpp/src/grid/bool_conversion.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "ATen/ATen.h" 3 | 4 | namespace fluid { 5 | 6 | bool toBool(const at::Tensor & self); 7 | 8 | } // namespace fluid 9 | -------------------------------------------------------------------------------- /solver_cpp/src/grid/cell_type.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | namespace fluid { 4 | // These are the same enum values used in Manta. We can't include grid.h 5 | // from Manta without pulling in the entire library, so we'll just redefine 6 | // them here. 7 | enum CellType { 8 | TypeNone = 0, 9 | TypeFluid = 1, 10 | TypeObstacle = 2, 11 | TypeEmpty = 4, 12 | TypeInflow = 8, 13 | TypeOutflow = 16, 14 | TypeOpen = 32, 15 | TypeStick = 128, 16 | TypeReserved = 256, 17 | TypeZeroPressure = (1<<15) 18 | }; 19 | 20 | } // namespace fluid 21 | 22 | -------------------------------------------------------------------------------- /solver_cpp/src/grid/grid.h: -------------------------------------------------------------------------------- 1 | #include "ATen/ATen.h" 2 | 3 | namespace fluid { 4 | 5 | typedef at::Tensor T; 6 | 7 | float getDx(at::Tensor self); 8 | 9 | T interpol(const T& self, const T& pos); 10 | 11 | void interpol1DWithFluid( 12 | const T& val_a, const T& is_fluid_a, 13 | const T& val_b, const T& is_fluid_b, 14 | const T& t_a, const T& t_b, 15 | T& is_fluid_ab, T& val_ab); 16 | 17 | T interpolWithFluid(const T& self, const T& flags, const T& pos); 18 | 19 | T getCentered(const T& self); 20 | 21 | T getAtMACX(const T& self); 22 | T getAtMACY(const T& self); 23 | T getAtMACZ(const T& self); 24 | 25 | T interpolComponent(const T& self, const T& pos, int c); 26 | 27 | T curl(const T& self); 28 | 29 | } // namespace fluid 30 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/.gitignore: -------------------------------------------------------------------------------- 1 | old 2 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/div.cpp: -------------------------------------------------------------------------------- 1 | #include "div.h" 2 | 3 | namespace fluid { 4 | 5 | // ***************************************************************************** 6 | // velocityDivergenceForward 7 | // ***************************************************************************** 8 | 9 | // Calculate the velocity divergence (with boundary cond modifications). This is 10 | // essentially a replica of makeRhs in Manta and FluidNet. 11 | // 12 | // input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 13 | // input flags - input occupancy grid 14 | // input UDiv - output divergence (scalar field). 15 | 16 | void velocityDivergenceForward(T& U, T& flags, T& UDiv) { 17 | // Check sizes 18 | AT_ASSERT(U.dim() == 5 && flags.dim() == 5 && UDiv.dim() == 5, 19 | "Dimension mismatch"); 20 | AT_ASSERT(flags.size(1) == 1, "flags is not scalar"); 21 | int bsz = flags.size(0); 22 | int d = flags.size(2); 23 | int h = flags.size(3); 24 | int w = flags.size(4); 25 | int bnd = 1; // Boundary width (hard coded) 26 | 27 | int z = 2; 28 | int y = 3; 29 | int x = 4; 30 | 31 | bool is_3d = (U.size(1) == 3); 32 | if (!is_3d) { 33 | AT_ASSERT(d == 1, "2D velocity field but zdepth > 1"); 34 | AT_ASSERT(U.size(1) == 2, "2D velocity field must have only 2 channels"); 35 | } 36 | AT_ASSERT((U.size(0) == bsz && U.size(2) == d && 37 | U.size(3) == h && U.size(4) == w), "Size mismatch"); 38 | AT_ASSERT(UDiv.is_same_size(flags), "Size mismatch"); 39 | 40 | AT_ASSERT(U.is_contiguous() && flags.is_contiguous() && 41 | UDiv.is_contiguous(), "Input is not contiguous"); 42 | 43 | T Uijk; // Velocity in ijk 44 | T Uijk_p; // Velocity in (i+1),(j+1),(k+1) 45 | 46 | // Remove the borders in x, y and z and build the i+1, j+1, k+1 tensor 47 | if (!is_3d) { 48 | Uijk = U.narrow(x, 1, w-2).narrow(y, 1, h-2); 49 | Uijk_p = Uijk.clone(); 50 | Uijk_p.select(1,0) = U.narrow(x, 2, w-2).narrow(y, 1, h-2).select(1,0); 51 | Uijk_p.select(1,1) = U.narrow(x, 1, w-2).narrow(y, 2, h-2).select(1,1); 52 | } else { 53 | Uijk = U.narrow(x, 1, w-2).narrow(y, 1, h-2).narrow(z, 1, d-2); 54 | Uijk_p = Uijk.clone(); 55 | Uijk_p.select(1,0) = U.narrow(x, 2, w-2).narrow(y, 1, h-2).narrow(z, 1, d-2).select(1,0); 56 | Uijk_p.select(1,1) = U.narrow(x, 1, w-2).narrow(y, 2, h-2).narrow(z, 1, d-2).select(1,1); 57 | Uijk_p.select(1,2) = U.narrow(x, 1, w-2).narrow(y, 1, h-2).narrow(z, 2, d-2).select(1,2); 58 | } 59 | 60 | // -div = u(i+1,j,k) - u(i,j,k) + 61 | // v(i,j+1,k) - v(i,j,k) + 62 | // w(i,j,k+1) - w(i,j,k) 63 | T div = Uijk.select(1,0) - Uijk_p.select(1,0) + 64 | Uijk.select(1,1) - Uijk_p.select(1,1); 65 | 66 | if (is_3d) { 67 | div += Uijk.select(1,2) - Uijk_p.select(1,2); 68 | } 69 | 70 | if (!is_3d) { 71 | UDiv.narrow(x, 1, w-2).narrow(y, 1, h-2) = div.view({bsz, 1, d, h-2, w-2}); 72 | } else { 73 | UDiv.narrow(x, 1, w-2).narrow(y, 1, h-2).narrow(z, 1, d-2) = div.view({bsz, 1, d-2, h-2, w-2}); 74 | } 75 | 76 | //Set div to 0 in obstacles 77 | T mask_obst = flags.eq(TypeObstacle); 78 | UDiv.masked_fill_(mask_obst, 0); 79 | } 80 | 81 | } // namespace fluid 82 | 83 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/div.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "ATen/ATen.h" 4 | #include "grid/cell_type.h" 5 | 6 | namespace fluid { 7 | 8 | typedef at::Tensor T; 9 | 10 | // ***************************************************************************** 11 | // velocityDivergenceForward 12 | // ***************************************************************************** 13 | 14 | // Calculate the velocity divergence (with boundary cond modifications). This is 15 | // essentially a replica of makeRhs in Manta and FluidNet. 16 | // 17 | // input U - input vel field (size(2) can be 2 or 3, indicating 2D / 3D) 18 | // input flags - input occupancy grid 19 | // input UDiv - output divergence (scalar field). 20 | 21 | void velocityDivergenceForward(T& U, T& flags, T& UDiv); 22 | 23 | } // namespace fluid 24 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/solve_linear_sys.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include "solve_linear_sys.h" 3 | 4 | namespace fluid { 5 | 6 | // ***************************************************************************** 7 | // solveLinearSystemJacobi 8 | // ***************************************************************************** 9 | 10 | // Solve the linear system using the Jacobi method. 11 | // Note: Since we don't receive a velocity field, we need to receive the is3D 12 | // flag from the caller. 13 | // 14 | // input p: The output pressure field (i.e. the solution to A * p = div) 15 | // input flags: The input flag grid. 16 | // input div: The velocity divergence. 17 | // input is3D: If true then we expect a 3D domain. 18 | // input pTol: OPTIONAL (default = 1e-5), ||p - p_prev|| termination cond. 19 | // input maxIter: OPTIONAL (default = 1000), max number of Jacobi iterations. 20 | // input verbose: OPTIONAL (default = false), if true print out iteration res. 21 | // 22 | // output: the max pTol across the batches. 23 | 24 | float solveLinearSystemJacobi 25 | ( 26 | T& p, 27 | T& flags, 28 | T& div, 29 | const bool is3D, 30 | const float p_tol = 1e-5, 31 | const int max_iter = 1000, 32 | const bool verbose = false 33 | ) { 34 | // Check arguments. 35 | AT_ASSERT(p.dim() == 5 && flags.dim() == 5 && div.dim() == 5, 36 | "Dimension mismatch"); 37 | AT_ASSERT(flags.size(1) == 1, "flags is not scalar"); 38 | int bsz = flags.size(0); 39 | int d = flags.size(2); 40 | int h = flags.size(3); 41 | int w = flags.size(4); 42 | int numel = d * h * w; 43 | AT_ASSERT(p.is_same_size(flags), "size mismatch"); 44 | AT_ASSERT(div.is_same_size(flags), "size mismatch"); 45 | if (!is3D) { 46 | AT_ASSERT(d == 1, "d > 1 for a 2D domain"); 47 | } 48 | 49 | AT_ASSERT(p.is_contiguous() && flags.is_contiguous() && 50 | div.is_contiguous(), "Input is not contiguous"); 51 | 52 | T p_prev = infer_type(p).zeros({bsz, 1, d, h, w}); 53 | T p_delta = infer_type(p).zeros({bsz, 1, d, h, w}); 54 | T p_delta_norm = infer_type(p).zeros({bsz}); 55 | 56 | if (max_iter < 1) { 57 | AT_ERROR("At least 1 iteration is needed (maxIter < 1)"); 58 | } 59 | 60 | // Initialize the pressure to zero. 61 | p.zero_(); 62 | 63 | // Start with the output of the next iteration going to pressure. 64 | T* cur_p = &p; 65 | T* cur_p_prev = &p_prev; 66 | //RealGrid* cur_pressure_prev = &pressure_prev; 67 | 68 | T residual; 69 | // T at_zero = infer_type(tensor_p).scalarTensor(0); 70 | 71 | int64_t iter = 0; 72 | while (true) { 73 | const int32_t bnd =1; 74 | // Kernel: Jacobi Iteration 75 | T mCont = infer_type(flags).ones({bsz, 1, d, h, w}).toType(at::kByte); // Continue mask 76 | 77 | T idx_x = infer_type(flags).arange(0, w).view({1,w}).expand({bsz, d, h, w}).toType(at::kLong); 78 | T idx_y = infer_type(idx_x).arange(0, h).view({1,h,1}).expand({bsz, d, h, w}); 79 | T idx_z = zeros_like(idx_x); 80 | if (is3D) { 81 | idx_z = infer_type(idx_x).arange(0, d).view({1,d,1,1}).expand({bsz, d, h, w}); 82 | } 83 | 84 | T idx_b = infer_type(flags).arange(0, bsz).view({bsz,1,1,1}).toType(at::kLong); 85 | idx_b = idx_b.expand({bsz,d,h,w}); 86 | 87 | T maskBorder = (idx_x < bnd).__or__ 88 | (idx_x > w - 1 - bnd).__or__ 89 | (idx_y < bnd).__or__ 90 | (idx_y > h - 1 - bnd); 91 | if (is3D) { 92 | maskBorder = maskBorder.__or__(idx_z < bnd).__or__ 93 | (idx_z > d - 1 - bnd); 94 | } 95 | maskBorder.unsqueeze_(1); 96 | 97 | // Zero pressure on the border. 98 | cur_p->masked_fill_(maskBorder, 0); 99 | mCont.masked_fill_(maskBorder, 0); 100 | 101 | T maskObstacle = flags.eq(TypeObstacle).__and__(mCont); 102 | cur_p->masked_fill_(maskObstacle, 0); 103 | mCont.masked_fill_(maskObstacle, 0); 104 | 105 | 106 | 107 | T zero_f = at::zeros_like(p); // Floating zero 108 | T zero_l = at::zeros_like(p).toType(at::kLong); // Long zero (for index) 109 | T zeroBy = at::zeros_like(p).toType(at::kByte); // Long zero (for index) 110 | // Otherwise, we are in a fluid or empty cell. 111 | // First, we get all the neighbors. 112 | 113 | T pC = *cur_p_prev; 114 | 115 | T i_l = zero_l.where( (idx_x <=0), idx_x - 1); 116 | T p1 = zero_f. 117 | where(mCont.ne(1), (*cur_p_prev).index({idx_b, zero_l, idx_z, idx_y, i_l}) 118 | .unsqueeze(1)); 119 | 120 | T i_r = zero_l.where( (idx_x > w - 1 - bnd), idx_x + 1); 121 | T p2 = zero_f. 122 | where(idx_x >= (w - 3 - bnd), (*cur_p_prev).index({idx_b, zero_l, idx_z, idx_y, i_r}) 123 | .unsqueeze(1)); 124 | 125 | T j_l = zero_l.where( (idx_y <= 0), idx_y - 1); 126 | T p3 = zero_f. 127 | where(mCont.ne(1), (*cur_p_prev).index({idx_b, zero_l, idx_z, j_l, idx_x}) 128 | .unsqueeze(1)); 129 | T j_r = zero_l.where( (idx_y > h - 1 - bnd), idx_y + 1); 130 | T p4 = zero_f. 131 | where(mCont.ne(1), (*cur_p_prev).index({idx_b, zero_l, idx_z, j_r, idx_x}) 132 | .unsqueeze(1)); 133 | 134 | T k_l = zero_l.where( (idx_z <= 0), idx_z - 1); 135 | T p5 = is3D ? zero_f. 136 | where(mCont.ne(1), (*cur_p_prev).index({idx_b, zero_l, k_l, idx_y, idx_x}) 137 | .unsqueeze(1)) : zero_f; 138 | T k_r = zero_l.where( (idx_z > d - 1 - bnd), idx_z + 1); 139 | T p6 = is3D ? zero_f. 140 | where(mCont.ne(1), (*cur_p_prev).index({idx_b, zero_l, k_r, idx_y, idx_x}) 141 | .unsqueeze(1)) : zero_f; 142 | 143 | T neighborLeftObs = mCont.__and__(zeroBy. 144 | where(mCont.ne(1), flags.index({idx_b, zero_l, idx_z, idx_y, i_l}).eq(TypeObstacle)).unsqueeze(1)); 145 | T neighborRightObs = mCont.__and__(zeroBy. 146 | where(mCont.ne(1), flags.index({idx_b, zero_l, idx_z, idx_y, i_r}).eq(TypeObstacle)).unsqueeze(1)); 147 | T neighborBotObs = mCont.__and__(zeroBy. 148 | where(mCont.ne(1), flags.index({idx_b, zero_l, idx_z, j_l, idx_x}).eq(TypeObstacle)).unsqueeze(1)); 149 | T neighborUpObs = mCont.__and__(zeroBy. 150 | where(mCont.ne(1), flags.index({idx_b, zero_l, idx_z, j_r, idx_x}).eq(TypeObstacle)).unsqueeze(1)); 151 | T neighborBackObs = zeroBy; 152 | T neighborFrontObs = zeroBy; 153 | 154 | if (is3D) { 155 | T neighborBackObs = mCont.__and__(zeroBy. 156 | where(mCont.ne(1), flags.index({idx_b, zero_l, k_l, idx_y, idx_x}).eq(TypeObstacle)).unsqueeze(1)); 157 | T neighborFrontObs = mCont.__and__(zeroBy. 158 | where(mCont.ne(1), flags.index({idx_b, zero_l, k_r, idx_y, idx_x}).eq(TypeObstacle)).unsqueeze(1)); 159 | } 160 | 161 | p1.masked_scatter_(neighborLeftObs, pC.masked_select(neighborLeftObs)); 162 | p2.masked_scatter_(neighborRightObs, pC.masked_select(neighborRightObs)); 163 | p3.masked_scatter_(neighborBotObs, pC.masked_select(neighborBotObs)); 164 | p4.masked_scatter_(neighborUpObs, pC.masked_select(neighborUpObs)); 165 | p5.masked_scatter_(neighborBackObs, pC.masked_select(neighborBackObs)); 166 | p6.masked_scatter_(neighborFrontObs, pC.masked_select(neighborFrontObs)); 167 | 168 | const float denom = is3D ? 6 : 4; 169 | (*cur_p).masked_scatter_(mCont, ((p1 + p2 + p3 + p4 + p5 + p6 + div) / denom) 170 | .masked_select(mCont)); 171 | 172 | // Currrent iteration output is now in cur_pressure 173 | 174 | // Now calculate the change in pressure up to a sign (the sign might be 175 | // incorrect, but we don't care). 176 | // p_delta = p - p_prev 177 | at::sub_out(p_delta, p, p_prev); 178 | p_delta.resize_({bsz, numel}); 179 | // Calculate L2 norm over dim 2. 180 | at::norm_out(p_delta_norm, p_delta, at::Scalar(2), 1); 181 | p_delta.resize_({bsz, 1, d, h, w}); 182 | residual = p_delta_norm.max(); 183 | if (verbose) { 184 | std::cout << "Jacobi iteration " << (iter + 1) << ": residual " 185 | << residual << std::endl; 186 | } 187 | 188 | if (at::Scalar(residual).toFloat() < p_tol) { 189 | if (verbose) { 190 | std::cout << "Jacobi max residual fell below p_tol (" << p_tol 191 | << ") (terminating)" << std::endl; 192 | } 193 | break; 194 | } 195 | 196 | iter++; 197 | if (iter >= max_iter) { 198 | if (verbose) { 199 | std::cout << "Jacobi max iteration count (" << max_iter 200 | << ") reached (terminating)" << std::endl; 201 | } 202 | break; 203 | } 204 | 205 | // We haven't yet terminated. 206 | auto tmp = cur_p; 207 | cur_p = cur_p_prev; 208 | cur_p_prev = tmp; 209 | } // end while 210 | 211 | // If we terminated with the cur_pressure pointing to the tmp array, then we 212 | // have to copy the pressure back into the output tensor. 213 | if (cur_p == &p_prev) { 214 | p.copy_(p_prev); // p = p_prev 215 | } 216 | 217 | // TODO: write mean-subtraction (FluidNet does it in Lua) 218 | return at::Scalar(residual).toFloat(); 219 | } 220 | 221 | } // namespace fluid 222 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/solve_linear_sys.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "grid/grid.h" 4 | #include "grid/cell_type.h" 5 | 6 | namespace fluid { 7 | 8 | // Solve the linear system using the Jacobi method. 9 | // Note: Since we don't receive a velocity field, we need to receive the is3D 10 | // flag from the caller. 11 | // 12 | // input p: The output pressure field (i.e. the solution to A * p = div) 13 | // input flags: The input flag grid. 14 | // input div: The velocity divergence. 15 | // input is3D: If true then we expect a 3D domain. 16 | // input pTol: OPTIONAL (default = 1e-5), ||p - p_prev|| termination cond. 17 | // input maxIter: OPTIONAL (default = 1000), max number of Jacobi iterations. 18 | // input verbose: OPTIONAL (default = false), if true print out iteration res. 19 | // 20 | // output: the max pTol across the batches. 21 | 22 | float solveLinearSystemJacobi 23 | ( 24 | T& p, 25 | T& flags, 26 | T& div, 27 | const bool is_3d, 28 | const float p_tol, 29 | const int max_iter, 30 | const bool verbose 31 | ); 32 | 33 | } // namespace fluid 34 | 35 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/update_vel.cpp: -------------------------------------------------------------------------------- 1 | #include "update_vel.h" 2 | 3 | namespace fluid { 4 | 5 | // ***************************************************************************** 6 | // velocityUpdateForward 7 | // ***************************************************************************** 8 | 9 | // Calculate the pressure gradient and subtract it into (i.e. calculate 10 | // U' = U - grad(p)). Some care must be taken with handling boundary conditions. 11 | // This function mimics correctVelocity in Manta. 12 | // NOTE: velocity update is done IN-PLACE. 13 | // 14 | // input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 15 | // input flags - input occupancy grid 16 | // input p - scalar pressure field. 17 | 18 | void velocityUpdateForward 19 | ( 20 | T& U, 21 | T& flags, 22 | T& pressure 23 | ) { 24 | // Check arguments. 25 | AT_ASSERT(U.dim() == 5 && flags.dim() == 5 && pressure.dim() == 5, 26 | "Dimension mismatch"); 27 | AT_ASSERT(flags.size(1) == 1, "flags is not scalar"); 28 | int b = flags.size(0); 29 | int d = flags.size(2); 30 | int h = flags.size(3); 31 | int w = flags.size(4); 32 | 33 | bool is3D = (U.size(1) == 3); 34 | if (!is3D) { 35 | AT_ASSERT(d == 1, "d > 1 for a 2D domain"); 36 | AT_ASSERT(U.size(4) == w, "2D velocity field must have only 2 channels"); 37 | } 38 | 39 | AT_ASSERT(U.size(0) == b && U.size(2) == d && U.size(3) == h 40 | && U.size(4) == w, "size mismatch"); 41 | AT_ASSERT(pressure.is_same_size(flags), "size mismatch"); 42 | 43 | AT_ASSERT(U.is_contiguous() && flags.is_contiguous() && 44 | pressure.is_contiguous(), "Input is not contiguous"); 45 | 46 | // First, we build the mask for detecting fluid cells. Borders are left untouched. 47 | T mask_fluid; // Fluid cells. 48 | T mask_fluid_i; // Fluid cells with (i-1) neighbour also a fluid. 49 | T mask_fluid_j; // Fluid cells with (j-1) neighbour also a fluid. 50 | T mask_fluid_k; // FLuid cells with (k-1) neighbour also a fluid. 51 | 52 | if (!is3D) { 53 | mask_fluid = flags.narrow(4, 1, w-2).narrow(3, 1, h-2).eq(fluid::TypeFluid); 54 | mask_fluid_i = mask_fluid.__and__ 55 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).eq(fluid::TypeFluid)); 56 | mask_fluid_j = mask_fluid.__and__ 57 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).eq(fluid::TypeFluid)); 58 | } else { 59 | mask_fluid = flags.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2).eq(fluid::TypeFluid); 60 | mask_fluid_i = mask_fluid.__and__ 61 | (flags.narrow(4, 0, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2).eq(fluid::TypeFluid)); 62 | mask_fluid_j = mask_fluid.__and__ 63 | (flags.narrow(4, 1, w-2).narrow(3, 0, h-2).narrow(2, 1, d-2).eq(fluid::TypeFluid)); 64 | mask_fluid_k = mask_fluid.__and__ 65 | (flags.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 0, d-2).eq(fluid::TypeFluid)); 66 | } 67 | 68 | // Cast into float or double tensor and cat into a single mask along chan. 69 | T mask_fluid_i_f = mask_fluid_i.type().toScalarType(U.type().scalarType()) 70 | .copy(mask_fluid_i); 71 | T mask_fluid_j_f = mask_fluid_j.type().toScalarType(U.type().scalarType()) 72 | .copy(mask_fluid_j); 73 | T mask_fluid_k_f; 74 | if (is3D) { 75 | mask_fluid_k_f = mask_fluid_k.type().toScalarType(U.type().scalarType()) 76 | .copy(mask_fluid_k); 77 | } 78 | 79 | T mask; 80 | if(!is3D) { 81 | mask = at::cat({mask_fluid_i_f, mask_fluid_j_f}, 1).contiguous(); 82 | } else { 83 | mask = at::cat({mask_fluid_i_f, mask_fluid_j_f, mask_fluid_k_f}, 1).contiguous(); 84 | } 85 | 86 | // pressure tensor. 87 | T Pijk; // Pressure at (i,j,k) in 3 channels (2 for 2D). 88 | T Pijk_m; // Pressure at chan 0: (i-1, j, k) 89 | // chan 1: (i, j-1, k) 90 | // chan 2: (i, j, k-1) 91 | 92 | if (!is3D) { 93 | Pijk = pressure.narrow(4, 1, w-2).narrow(3, 1, h-2); 94 | Pijk = Pijk.clone().expand({b, 2, d, h-2, w-2}); 95 | Pijk_m = Pijk.clone().expand({b, 2, d, h-2, w-2}); 96 | Pijk_m.select(1,0) = pressure.narrow(4, 0, w-2).narrow(3, 1, h-2).squeeze(1); 97 | Pijk_m.select(1,1) = pressure.narrow(4, 1, w-2).narrow(3, 0, h-2).squeeze(1); 98 | } else { 99 | Pijk = pressure.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2); 100 | Pijk = Pijk.clone().expand({b, 3, d-2, h-2, w-2}); 101 | Pijk_m = Pijk.clone().expand({b, 3, d-2, h-2, w-2}); 102 | Pijk_m.select(1,0) = pressure.narrow(4, 0, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2).squeeze(1); 103 | Pijk_m.select(1,1) = pressure.narrow(4, 1, w-2).narrow(3, 0, h-2).narrow(2, 1, d-2).squeeze(1); 104 | Pijk_m.select(1,2) = pressure.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 0, d-2).squeeze(1); 105 | } 106 | 107 | // u = u - grad(p) 108 | // grad(p) = [[ p(i,j,k) - p(i-1,j,k) ] 109 | // [ p(i,j,k) - p(i,j-1,k) ] 110 | // [ p(i,j,k) - p(i,j,k-1) ]] 111 | if (!is3D) { 112 | U.narrow(4, 1, w-2).narrow(3, 1, h-2) = mask * 113 | (U.narrow(4, 1, w-2).narrow(3, 1, h-2) - (Pijk - Pijk_m)); 114 | } else { 115 | U.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2) = mask * 116 | (U.narrow(4, 1, w-2).narrow(3, 1, h-2).narrow(2, 1, d-2) - (Pijk - Pijk_m)); 117 | } 118 | 119 | } 120 | 121 | } // namespace fluid 122 | -------------------------------------------------------------------------------- /solver_cpp/src/projection/update_vel.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "ATen/ATen.h" 4 | #include "grid/cell_type.h" 5 | 6 | namespace fluid { 7 | 8 | typedef at::Tensor T; 9 | 10 | // ***************************************************************************** 11 | // velocityUpdateForward 12 | // ***************************************************************************** 13 | 14 | // Calculate the pressure gradient and subtract it into (i.e. calculate 15 | // U' = U - grad(p)). Some care must be taken with handling boundary conditions. 16 | // This function mimics correctVelocity in Manta. 17 | // NOTE: velocity update is done IN-PLACE. 18 | // 19 | // input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 20 | // input flags - input occupancy grid 21 | // input pressure - scalar pressure field. 22 | 23 | void velocityUpdateForward(T& U, T& flags, T& pressure); 24 | 25 | } // namespace fluid 26 | -------------------------------------------------------------------------------- /solver_cpp/src/sourceTerms/.gitignore: -------------------------------------------------------------------------------- 1 | old 2 | -------------------------------------------------------------------------------- /solver_cpp/src/sourceTerms/source_term.cpp: -------------------------------------------------------------------------------- 1 | #include "source_term.h" 2 | 3 | namespace fluid { 4 | 5 | typedef at::Tensor T; 6 | 7 | // ***************************************************************************** 8 | // addBuoyancy 9 | // ***************************************************************************** 10 | 11 | // Add buoyancy force. AddBuoyancy has a dt term. 12 | // Note: Buoyancy is added IN-PLACE. 13 | // 14 | // @input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 15 | // @input flags - input occupancy grid 16 | // @input density - scalar density grid. 17 | // @input gravity - 3D vector indicating direction of gravity. 18 | // @input dt - scalar timestep. 19 | 20 | void addBuoyancy 21 | ( 22 | T& U, T& flags, T& density, T& gravity, 23 | const float dt 24 | ) { 25 | // Argument check 26 | AT_ASSERT(U.dim() == 5 && flags.dim() == 5 && density.dim() == 5, 27 | "Dimension mismatch"); 28 | AT_ASSERT(flags.size(1) == 1, "flags is not scalar"); 29 | int bsz = flags.size(0); 30 | int d = flags.size(2); 31 | int h = flags.size(3); 32 | int w = flags.size(4); 33 | 34 | bool is3D = (U.size(1) == 3); 35 | 36 | int bnd = 1; 37 | 38 | if (!is3D) { 39 | AT_ASSERT(d == 1, "2D velocity field but zdepth > 1"); 40 | AT_ASSERT(U.size(1) == 2, "2D velocity field must have only 2 channels"); 41 | } 42 | AT_ASSERT((U.size(0) == bsz && U.size(2) == d && 43 | U.size(3) == h && U.size(4) == w), "Size mismatch"); 44 | AT_ASSERT(density.is_same_size(flags), "Size mismatch"); 45 | 46 | AT_ASSERT(U.is_contiguous() && flags.is_contiguous() && 47 | density.is_contiguous(), "Input is not contiguous"); 48 | 49 | AT_ASSERT(gravity.dim() == 1 && gravity.size(0) == 3, 50 | "Gravity must be a 3D vector (even in 2D)"); 51 | 52 | T strength = - gravity * (dt / getDx(flags)); 53 | 54 | T i = infer_type(flags).arange(0, w).view({1,w}).expand({bsz, d, h, w}).toType(at::kLong); 55 | T j = infer_type(i).arange(0, h).view({1,h,1}).expand({bsz, d, h, w}); 56 | T k = zeros_like(i); 57 | if (is3D) { 58 | k = infer_type(i).arange(0, d).view({1,d,1,1}).expand({bsz, d, h, w}); 59 | } 60 | T zero = zeros_like(i); 61 | T zeroBy = zero.toType(at::kByte); 62 | T zero_f = zero.toType(infer_type(density)); 63 | 64 | T idx_b = infer_type(i).arange(0, bsz).view({bsz,1,1,1}); 65 | idx_b = idx_b.expand({bsz,d,h,w}); 66 | 67 | T maskBorder = (i < bnd).__or__ 68 | (i > w - 1 - bnd).__or__ 69 | (j < bnd).__or__ 70 | (j > h - 1 - bnd); 71 | if (is3D) { 72 | maskBorder = maskBorder.__or__(k < bnd).__or__ 73 | (k > d - 1 - bnd); 74 | } 75 | maskBorder = maskBorder.unsqueeze(1); 76 | 77 | // No buoyancy on the border. Set continue (mCont) to false. 78 | T mCont = ones_like(zeroBy).unsqueeze(1); 79 | mCont.masked_fill_(maskBorder, 0); 80 | 81 | T isFluid = flags.eq(TypeFluid).__and__(mCont); 82 | mCont.masked_fill_(isFluid.ne(1), 0); 83 | mCont.squeeze_(1); 84 | 85 | T fluid100 = zeroBy.where( i <= 0, (flags.index({idx_b, zero, k, j, i-1}).eq(TypeFluid))).__and__(mCont); 86 | T factor = 0.5 * strength[0] * (density.squeeze(1) + 87 | zero_f.where(i <= 0, density.index({idx_b, zero, k, j, i-1})) ); 88 | U.select(1,0).masked_scatter_(fluid100, (U.select(1,0) + factor).masked_select(fluid100)); 89 | 90 | T fluid010 = zeroBy.where( j <= 0, (flags.index({idx_b, zero, k, j-1, i}).eq(TypeFluid))).__and__(mCont); 91 | factor = 0.5 * strength[1] * (density.squeeze(1) + 92 | zero_f.where( j <= 0, density.index({idx_b, zero, k, j-1, i})) ); 93 | U.select(1,1).masked_scatter_(fluid010, (U.select(1,1) + factor).masked_select(fluid010)); 94 | 95 | if (is3D) { 96 | T fluid001 = zeroBy.where( j <= 0, (flags.index({idx_b, zero, k-1, j, i}).eq(TypeFluid))).__and__(mCont); 97 | factor = 0.5 * strength[2] * (density.squeeze(1) + 98 | zero_f.where(k <= 1, density.index({idx_b, zero, k-1, j, i})) ); 99 | U.select(1,2).masked_scatter_(fluid001, (U.select(1,2) + factor).masked_select(fluid001)); 100 | 101 | } 102 | 103 | } 104 | 105 | // ***************************************************************************** 106 | // addGravity 107 | // ***************************************************************************** 108 | 109 | // Add gravity force. It has a dt term. 110 | // Note: gravity is added IN-PLACE. 111 | // 112 | // @input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 113 | // @input flags - input occupancy grid 114 | // @input gravity - 3D vector indicating direction of gravity. 115 | // @input dt - scalar timestep. 116 | 117 | void addGravity 118 | ( 119 | T& U, T& flags, T& gravity, 120 | const float dt 121 | ) { 122 | // Argument check 123 | AT_ASSERT(U.dim() == 5 && flags.dim() == 5, "Dimension mismatch"); 124 | AT_ASSERT(flags.size(1) == 1, "flags is not scalar"); 125 | int bsz = flags.size(0); 126 | int d = flags.size(2); 127 | int h = flags.size(3); 128 | int w = flags.size(4); 129 | 130 | bool is3D = (U.size(1) == 3); 131 | 132 | int bnd = 1; 133 | if (!is3D) { 134 | AT_ASSERT(d == 1, "2D velocity field but zdepth > 1"); 135 | AT_ASSERT(U.size(1) == 2, "2D velocity field must have only 2 channels"); 136 | } 137 | AT_ASSERT((U.size(0) == bsz && U.size(2) == d && 138 | U.size(3) == h && U.size(4) == w), "Size mismatch"); 139 | 140 | AT_ASSERT(U.is_contiguous() && flags.is_contiguous(), "Input is not contiguous"); 141 | 142 | AT_ASSERT(gravity.dim() == 1 && gravity.size(0) == 3, 143 | "Gravity must be a 3D vector (even in 2D)"); 144 | 145 | T force = gravity * (dt / getDx(flags)); 146 | 147 | T i = infer_type(flags).arange(0, w).view({1,w}).expand({bsz, d, h, w}).toType(at::kLong); 148 | T j = infer_type(i).arange(0, h).view({1,h,1}).expand({bsz, d, h, w}); 149 | T k = zeros_like(i); 150 | if (is3D) { 151 | k = infer_type(i).arange(0, d).view({1,d,1,1}).expand({bsz, d, h, w}); 152 | } 153 | T zero = zeros_like(i); 154 | T zeroBy = zero.toType(at::kByte); 155 | T zero_f = zero.toType(infer_type(U)); 156 | 157 | T idx_b = infer_type(i).arange(0, bsz).view({bsz,1,1,1}); 158 | idx_b = idx_b.expand({bsz,d,h,w}); 159 | 160 | T maskBorder = (i < bnd).__or__ 161 | (i > w - 1 - bnd).__or__ 162 | (j < bnd).__or__ 163 | (j > h - 1 - bnd); 164 | if (is3D) { 165 | maskBorder = maskBorder.__or__(k < bnd).__or__ 166 | (k > d - 1 - bnd); 167 | } 168 | maskBorder = maskBorder.unsqueeze(1); 169 | 170 | // No buoyancy on the border. Set continue (mCont) to false. 171 | T mCont = ones_like(zeroBy).unsqueeze(1); 172 | mCont.masked_fill_(maskBorder, 0); 173 | 174 | T cur_fluid = flags.eq(TypeFluid).__and__(mCont); 175 | T cur_empty = flags.eq(TypeEmpty).__and__(mCont); 176 | 177 | T mNotFluidNotEmpt = cur_fluid.ne(1).__and__(cur_empty.ne(1)); 178 | mCont.masked_fill_(mNotFluidNotEmpt, 0); 179 | 180 | mCont.squeeze_(1); 181 | 182 | T fluid100 = (zeroBy.where( i <= 0, (flags.index({idx_b, zero, k, j, i-1}).eq(TypeFluid))) 183 | .__or__(( zeroBy.where( i <= 0, (flags.index({idx_b, zero, k, j, i-1}).eq(TypeEmpty)))) 184 | .__and__(cur_fluid.squeeze(1)))).__and__(mCont); 185 | U.select(1,0).masked_scatter_(fluid100, (U.select(1,0) + force[0]).masked_select(fluid100)); 186 | 187 | T fluid010 = (zeroBy.where( j <= 0, (flags.index({idx_b, zero, k, j-1, i}).eq(TypeFluid))) 188 | .__or__(( zeroBy.where( j <= 0, (flags.index({idx_b, zero, k, j-1, i}).eq(TypeEmpty)))) 189 | .__and__(cur_fluid.squeeze(1))) ).__and__(mCont); 190 | U.select(1,1).masked_scatter_(fluid010, (U.select(1,1) + force[1]).masked_select(fluid010)); 191 | 192 | if (is3D) { 193 | T fluid001 = (zeroBy.where( k <= 0, (flags.index({idx_b, zero, k-1, j, i}).eq(TypeFluid))) 194 | .__or__(( zeroBy.where( k <= 0, (flags.index({idx_b, zero, k-1, j, i}).eq(TypeEmpty)))) 195 | .__and__(cur_fluid.squeeze(1)))).__and__(mCont); 196 | U.select(1,2).masked_scatter_(fluid001, (U.select(1,2) + force[2]).masked_select(fluid001)); 197 | } 198 | 199 | } 200 | 201 | } // namespace fluid 202 | -------------------------------------------------------------------------------- /solver_cpp/src/sourceTerms/source_term.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "ATen/ATen.h" 4 | #include "grid/grid.h" 5 | #include "grid/cell_type.h" 6 | 7 | namespace fluid { 8 | 9 | typedef at::Tensor T; 10 | 11 | // ***************************************************************************** 12 | // addBuoyancy 13 | // ***************************************************************************** 14 | 15 | // Add buoyancy force. AddBuoyancy has a dt term. 16 | // Note: Buoyancy is added IN-PLACE. 17 | // 18 | // @input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 19 | // @input flags - input occupancy grid 20 | // @input density - scalar density grid. 21 | // @input gravity - 3D vector indicating direction of gravity. 22 | // @input dt - scalar timestep. 23 | 24 | void addBuoyancy(T& tensor_u, T& tensor_flags, T& tensor_density, 25 | T& tensor_gravity, const float dt); 26 | 27 | // ***************************************************************************** 28 | // addGravity 29 | // ***************************************************************************** 30 | 31 | // Add gravity force. It has a dt term. 32 | // Note: gravity is added IN-PLACE. 33 | // 34 | // @input U - vel field (size(2) can be 2 or 3, indicating 2D / 3D) 35 | // @input flags - input occupancy grid 36 | // @input gravity - 3D vector indicating direction of gravity. 37 | // @input dt - scalar timestep. 38 | 39 | void addGravity(T& tensor_u, T& tensor_flags, T& tensor_gravity, 40 | const float dt); 41 | 42 | } // namespace fluid 43 | -------------------------------------------------------------------------------- /solver_cpp/test/.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | -------------------------------------------------------------------------------- /solver_cpp/test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | SET(SIM_SRC test_fluid.cpp) 2 | 3 | ADD_EXECUTABLE(fluidnet_sim ${SIM_SRC}) 4 | 5 | TARGET_INCLUDE_DIRECTORIES(fluidnet_sim PUBLIC 6 | $ 7 | ) 8 | LINK_DIRECTORIES(fluidnet_sim $) 9 | 10 | TARGET_LINK_LIBRARIES(fluidnet_sim FluidNet ATen) 11 | TARGET_LINK_LIBRARIES(fluidnet_sim ${OpenCV_LIBS} ) 12 | -------------------------------------------------------------------------------- /solver_cpp/test/load_manta_data.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include // glob(), globfree() to parse through a directory 11 | #include 12 | #include //memset 13 | 14 | #include "ATen/ATen.h" 15 | #include "grid/bool_conversion.h" 16 | 17 | // Load data from one .bin file generated in python with Manta 18 | bool loadMantaFile 19 | ( 20 | std::string fname, 21 | at::Tensor& p, 22 | at::Tensor& U, 23 | at::Tensor& flags, 24 | at::Tensor& density, 25 | bool& is3D 26 | ) 27 | { 28 | // Load files in CPU only! 29 | auto && Tfloat = CPU(at::kFloat); 30 | auto && Tint = CPU(at::kInt); 31 | auto && Tundef = at::getType(at::Backend::Undefined, at::ScalarType::Undefined); 32 | 33 | if (p.type() != Tundef || U.type() != Tundef || density.type() 34 | != Tundef || flags.type() != Tundef) { 35 | AT_ERROR("Load Manta File: input tensors must be Undefined"); 36 | } 37 | 38 | std::fstream file; 39 | file.open(fname, std::ios::in | std::ios::binary); 40 | 41 | if(file.fail()) { 42 | std::cout << "Unable to open the data file!!!" <(&transpose), sizeof(int) ); 54 | file.read(reinterpret_cast(&nx), sizeof(int) ); 55 | file.read(reinterpret_cast(&ny), sizeof(int) ); 56 | file.read(reinterpret_cast(&nz), sizeof(int) ); 57 | file.read(reinterpret_cast(&is3D_), sizeof(int) ); 58 | is3D = (is3D_ == 1); 59 | const int numel = nx*ny*nz; 60 | float velx[numel]; 61 | float vely[numel]; 62 | float velz[numel]; 63 | 64 | // Clones are necessay everywhere after reading because 65 | // tensorFromBlob does not own memory (unable to resize). 66 | 67 | file.read(reinterpret_cast(&velx), sizeof(velx) ); 68 | file.read(reinterpret_cast(&vely), sizeof(vely) ); 69 | at::Tensor Ux_read = Tfloat.tensorFromBlob(velx, {numel}); 70 | at::Tensor Uy_read = Tfloat.tensorFromBlob(vely, {numel}); 71 | at::Tensor Ux = Ux_read.clone(); 72 | at::Tensor Uy = Uy_read.clone(); 73 | at::Tensor Uz; 74 | if (is3D) { 75 | // nz -= 1; 76 | //float velz[numel]; 77 | file.read(reinterpret_cast(&velz), sizeof(velz) ); 78 | at::Tensor Uz_read = Tfloat.tensorFromBlob(velz, {numel}); 79 | Uz = Uz_read.clone(); 80 | } 81 | 82 | float pres[numel]; 83 | file.read(reinterpret_cast(&pres), sizeof(pres) ); 84 | at::Tensor p_read = Tfloat.tensorFromBlob(pres, {numel}); 85 | p = p_read.clone(); 86 | 87 | int flagIN[numel]; 88 | file.read(reinterpret_cast(&flagIN), sizeof(flagIN) ); 89 | at::Tensor flags_read = Tint.tensorFromBlob(flagIN, {numel}); 90 | flags = flags_read.clone(); 91 | flags = flags.toType(Tfloat); 92 | 93 | float rho[numel]; 94 | file.read(reinterpret_cast(&rho), sizeof(rho) ); 95 | at::Tensor density_read = Tfloat.tensorFromBlob(rho, {numel}); 96 | density = density_read.clone(); 97 | 98 | Ux.resize_({1, 1, nz, ny, nx}); 99 | Uy.resize_({1, 1, nz, ny, nx}); 100 | if (is3D) { 101 | Uz.resize_({1, 1, nz, ny, nx}); 102 | } 103 | p.resize_({1, 1, nz, ny, nx}); 104 | flags.resize_({1, 1, nz, ny, nx}); 105 | density.resize_({1, 1, nz, ny, nx}); 106 | 107 | if (is3D) { 108 | U = at::cat({Ux, Uy, Uz}, 1).contiguous(); 109 | } 110 | else{ 111 | U = at::cat({Ux, Uy}, 1).contiguous(); 112 | } 113 | 114 | file.close(); 115 | return true; 116 | } 117 | 118 | // Parse through a directory and store file names in a vector of strings 119 | // https://stackoverflow.com/questions/8401777/simple-glob-in-c-on-unix-system 120 | 121 | std::vector globVector(const std::string& pattern){ 122 | glob_t glob_result; 123 | std::memset(&glob_result, 0, sizeof(glob_result)); 124 | 125 | int rtrn_val = glob(pattern.c_str(), GLOB_TILDE, NULL, &glob_result); 126 | if(rtrn_val != 0) { 127 | globfree(&glob_result); 128 | std::string ss; 129 | ss = "glob() failed with return value " + std::to_string(rtrn_val) + ". Check the file name " + pattern + "."; 130 | const char *css = ss.c_str(); 131 | AT_ERROR(css); 132 | } 133 | 134 | std::vector files; 135 | 136 | for(unsigned int i=0; i files = globVector(path); 160 | if (files.size() != 16){ 161 | std::string ss; 162 | ss = "loadMantaBatch(" + fn + ") must have 16 files per batch"; 163 | const char *css = ss.c_str(); 164 | AT_ERROR(css); 165 | } 166 | std::vector p; 167 | std::vector U; 168 | std::vector flags; 169 | std::vector density; 170 | 171 | for (auto const& file: files){ 172 | at::Tensor curP; 173 | at::Tensor curU; 174 | at::Tensor curFlags; 175 | at::Tensor curDensity; 176 | bool curIs3D; 177 | bool succes = loadMantaFile(file, curP, curU, curFlags, curDensity, curIs3D); 178 | p.push_back(curP); 179 | U.push_back(curU); 180 | flags.push_back(curFlags); 181 | density.push_back(curDensity); 182 | is3D = curIs3D; 183 | } 184 | 185 | p_out = at::cat(p, 0); 186 | U_out = at::cat(U, 0); 187 | flags_out = at::cat(flags, 0); 188 | density_out = at::cat (density, 0); 189 | } 190 | 191 | // Make sure we load samples from file that aren't all the same, otherwise 192 | // we're not really testing the batch dimension. 193 | void assertNotAllEqual(at::Tensor t) 194 | { 195 | if (t.size(0) == 1){ 196 | return; // Only one sample, so it's correct. 197 | } 198 | 199 | at::Tensor first = t.select(0,0); 200 | at::Tensor others = t.select(0,10); 201 | if(fluid::toBool(others - at::max(at::abs(first.expand_as(others))) < 1e-5)){ 202 | AT_ERROR("All samples are equal!"); 203 | } 204 | } 205 | -------------------------------------------------------------------------------- /solver_cpp/test/plot_utils.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #include "ATen/ATen.h" 6 | 7 | #include 8 | #include 9 | #include 10 | 11 | void plotTensor2D(at::Tensor Ten, int img_w, int img_h, std::string wname) 12 | { 13 | at::Tensor T_min = at::min(Ten); 14 | at::Tensor T_max = at::max(Ten); 15 | at::Tensor T_scaled = (Ten - T_min) / (T_max - T_min); 16 | 17 | 18 | int height = Ten.size(Ten.dim()-2); 19 | int width = Ten.size(Ten.dim()-1); 20 | 21 | cv::Mat img_tensor(width, height, CV_32FC1); 22 | cv::Mat img_rsz(img_w, img_h, CV_8UC1); 23 | cv::Mat img8uc; 24 | cv::Mat img_col; 25 | 26 | if (Ten.dim() == 5) { 27 | auto T_a = T_scaled.accessor(); 28 | for (int h = 0; h < height; h++) { 29 | for (int w = 0; w < width; w++) { 30 | 31 | img_tensor.at(w,h) = T_a[0][0][0][h][w]; 32 | } 33 | } } 34 | else if (Ten.dim() == 4) { 35 | auto T_a = T_scaled.accessor(); 36 | for (int h = 0; h < height; h++) { 37 | for (int w = 0; w < width; w++) { 38 | img_tensor.at(w,h) = T_a[0][0][h][w]; 39 | } 40 | } } 41 | else if (Ten.dim() == 3) { 42 | auto T_a = T_scaled.accessor(); 43 | for (int h = 0; h < height; h++) { 44 | for (int w = 0; w < width; w++) { 45 | img_tensor.at(w,h) = T_a[0][h][w]; 46 | } 47 | } } 48 | else if (Ten.dim() == 2) { 49 | auto T_a = T_scaled.accessor(); 50 | for (int h = 0; h < height; h++) { 51 | for (int w = 0; w < width; w++) { 52 | img_tensor.at(w,h) = T_a[h][w]; 53 | } 54 | } } 55 | else { 56 | AT_ERROR("Plot: error in input dimension!"); 57 | } 58 | 59 | double min; 60 | double max; 61 | cv::minMaxLoc(img_tensor, &min, &max); 62 | 63 | //CV_8UC1 conversion 64 | img_tensor.convertTo(img8uc,CV_8UC1, 255/(max-min),-255*min / (max - min)); 65 | resize(img8uc, img_rsz, img_rsz.size(), 0, 0, cv::INTER_NEAREST); 66 | applyColorMap(img_rsz, img_col, cv::COLORMAP_JET); 67 | std::string out_name = wname + ".png"; 68 | cv::imwrite(out_name, img_col); 69 | } 70 | 71 | -------------------------------------------------------------------------------- /solver_cpp/test/type_test.h: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #ifndef _MSC_VER 4 | # include 5 | #endif 6 | #include 7 | #include 8 | #include 9 | 10 | template 11 | std::string 12 | type_name() 13 | { 14 | typedef typename std::remove_reference::type TR; 15 | std::unique_ptr own 16 | ( 17 | #ifndef _MSC_VER 18 | abi::__cxa_demangle(typeid(TR).name(), nullptr, 19 | nullptr, nullptr), 20 | #else 21 | nullptr, 22 | #endif 23 | std::free 24 | ); 25 | std::string r = own != nullptr ? own.get() : typeid(TR).name(); 26 | if (std::is_const::value) 27 | r += " const"; 28 | if (std::is_volatile::value) 29 | r += " volatile"; 30 | if (std::is_lvalue_reference::value) 31 | r += "&"; 32 | else if (std::is_rvalue_reference::value) 33 | r += "&&"; 34 | return r; 35 | } 36 | 37 | -------------------------------------------------------------------------------- /trained_models/ScaleNet_ShortTerm_LongTermLoss/ScaleNet_ShortTerm_LongTermLoss_saved.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | 5 | from lib import fluid, MultiScaleNet 6 | from math import inf 7 | 8 | class _ScaleNet(nn.Module): 9 | def __init__(self, mconf): 10 | super(_ScaleNet, self).__init__() 11 | self.mconf = mconf 12 | 13 | def forward(self, x): 14 | bsz = x.size(0) 15 | # Rehaspe form (b x chan x d x h x w) to (b x -1) 16 | y = x.view(bsz, -1) 17 | # Calculate std using Bessel's correction (correction with n/n-1) 18 | std = torch.std(y, dim=1, keepdim=True) # output is size (b x 1) 19 | scale = torch.clamp(std, \ 20 | self.mconf['normalizeInputThreshold'] , inf) 21 | scale = scale.view(bsz, 1, 1, 1, 1) 22 | 23 | return scale 24 | 25 | class _HiddenConvBlock(nn.Module): 26 | def __init__(self, dropout=True): 27 | super(_HiddenConvBlock, self).__init__() 28 | layers = [ 29 | nn.Conv2d(16, 16, kernel_size=3, padding = 1), 30 | nn.ReLU(inplace=True), 31 | nn.Conv2d(16, 16, kernel_size=3, padding = 1), 32 | nn.ReLU(), 33 | ] 34 | if dropout: 35 | layers.append(nn.Dropout()) 36 | self.block = nn.Sequential(*layers) 37 | 38 | def forward(self, x): 39 | return self.block(x) 40 | 41 | 42 | class FluidNet(nn.Module): 43 | # For now, only 2D model. Add 2D/3D option. Only known from data! 44 | # Also, build model with MSE of pressure as loss func, therefore input is velocity 45 | # and output is pressure, to be compared to target pressure. 46 | def __init__(self, mconf, dropout=True): 47 | super(FluidNet, self).__init__() 48 | 49 | self.dropout = dropout 50 | self.mconf = mconf 51 | self.inDims = mconf['inputDim'] 52 | self.is3D = mconf['is3D'] 53 | 54 | self.scale = _ScaleNet(self.mconf) 55 | # Input channels = 3 (inDims, flags) 56 | # We add padding to make sure that Win = Wout and Hin = Hout with ker_size=3 57 | self.conv1 = torch.nn.Conv2d(self.inDims, 16, kernel_size=3, padding=1) 58 | 59 | self.modDown1 = torch.nn.AvgPool2d(kernel_size=2) 60 | self.modDown2 = torch.nn.AvgPool2d(kernel_size=4) 61 | 62 | self.convBank = _HiddenConvBlock(dropout) 63 | 64 | #self.upscale1 = torch.nn.Upsample(scale_factor=2, mode='nearest') 65 | #self.upscale2 = torch.nn.Upsample(scale_factor=4, mode='nearest') 66 | 67 | self.deconv1 = torch.nn.ConvTranspose2d(16, 16, kernel_size=2, stride=2) 68 | self.deconv2 = torch.nn.ConvTranspose2d(16, 16, kernel_size=4, stride=4) 69 | 70 | self.conv2 = torch.nn.Conv2d(16*3, 16, kernel_size=1) 71 | 72 | # Output channels = 1 (pressure) 73 | self.convOut = torch.nn.Conv2d(16, 1, kernel_size=1) 74 | 75 | # MultiScaleNet 76 | self.multiScale = MultiScaleNet(self.inDims) 77 | 78 | def forward(self, input_): 79 | 80 | # data indexes | | 81 | # (dim 1) | 2D | 3D 82 | # ---------------------------------------- 83 | # DATA: 84 | # pDiv | 0 | 0 85 | # UDiv | 1:3 | 1:4 86 | # flags | 3 | 4 87 | # densityDiv | 4 | 5 88 | # TARGET: 89 | # p | 0 | 0 90 | # U | 1:3 | 1:4 91 | # density | 3 | 4 92 | 93 | # For now, we work ONLY in 2d 94 | 95 | assert self.is3D == False, 'Input can only be 2D' 96 | 97 | assert self.mconf['inputChannels']['pDiv'] or \ 98 | self.mconf['inputChannels']['UDiv'] or \ 99 | self.mconf['inputChannels']['div'], 'Choose at least one field (U, div or p).' 100 | 101 | pDiv = None 102 | UDiv = None 103 | div = None 104 | 105 | # Flags are always loaded 106 | if self.is3D: 107 | flags = input_[:,4].unsqueeze(1) 108 | else: 109 | flags = input_[:,3].unsqueeze(1).contiguous() 110 | 111 | if (self.mconf['inputChannels']['pDiv'] or (self.mconf['normalizeInput'] \ 112 | and self.mconf['normalizeInputChan'] == 'pDiv')): 113 | pDiv = input_[:,0].unsqueeze(1).contiguous() 114 | 115 | if (self.mconf['inputChannels']['UDiv'] or self.mconf['inputChannels']['div'] \ 116 | or (self.mconf['normalizeInput'] \ 117 | and self.mconf['normalizeInputChan'] == 'UDiv')): 118 | if self.is3D: 119 | UDiv = input_[:,1:4].contiguous() 120 | else: 121 | UDiv = input_[:,1:3].contiguous() 122 | 123 | if 'periodic-x' in self.mconf and 'periodic-y' in self.mconf: 124 | U_temp = UDiv.clone() 125 | 126 | # Apply setWallBcs to zero out obstacles velocities on the boundary 127 | #UDiv = fluid.setWallBcs(UDiv, flags) 128 | if 'periodic-x' in self.mconf and 'periodic-y' in self.mconf: 129 | if self.mconf['periodic-x']: 130 | UDiv[:,1,:,:,1] = U_temp[:,1,:,:,UDiv.size(4)-1] 131 | if self.mconf['periodic-y']: 132 | UDiv[:,0,:,1] = U_temp[:,0,:,UDiv.size(3)-1] 133 | 134 | if self.mconf['inputChannels']['div']: 135 | div = fluid.velocityDivergence(UDiv, flags) 136 | 137 | # Apply scale to input 138 | if self.mconf['normalizeInput']: 139 | if self.mconf['normalizeInputChan'] == 'UDiv': 140 | s = self.scale(UDiv) 141 | elif self.mconf['normalizeInputChan'] == 'pDiv': 142 | s = self.scale(pDiv) 143 | elif self.mconf['normalizeInputChan'] == 'div': 144 | s = self.scale(div) 145 | else: 146 | raise Exception('Incorrect normalize input channel.') 147 | 148 | if pDiv is not None: 149 | pDiv = torch.div(pDiv, s) 150 | if UDiv is not None: 151 | UDiv = torch.div(UDiv, s) 152 | if div is not None: 153 | div = torch.div(div, s) 154 | 155 | x = torch.FloatTensor(input_.size(0), \ 156 | self.inDims, \ 157 | input_.size(2), \ 158 | input_.size(3), \ 159 | input_.size(4)).type_as(input_) 160 | 161 | chan = 0 162 | if self.mconf['inputChannels']['pDiv']: 163 | x[:, chan] = pDiv[:,0] 164 | chan += 1 165 | elif self.mconf['inputChannels']['UDiv']: 166 | if self.is3D: 167 | x[:,chan:(chan+3)] = UDiv 168 | chan += 3 169 | else: 170 | x[:,chan:(chan+2)] = UDiv 171 | chan += 2 172 | elif self.mconf['inputChannels']['div']: 173 | x[:, chan] = div[:,0] 174 | chan += 1 175 | 176 | # FlagsToOccupancy creates a [0,1] grid out of the manta flags 177 | x[:,chan,:,:,:] = fluid.flagsToOccupancy(flags).squeeze(1) 178 | 179 | if not self.is3D: 180 | # Squeeze unary dimension as we are in 2D 181 | x = torch.squeeze(x,2) 182 | 183 | if self.mconf['model'] == 'ScaleNet': 184 | p = self.multiScale(x) 185 | 186 | else: 187 | # Inital layers 188 | x = F.relu(self.conv1(x)) 189 | 190 | # We divide the network in 3 banks, applying average pooling 191 | x1 = self.modDown1(x) 192 | x2 = self.modDown2(x) 193 | 194 | # Process every bank in parallel 195 | x0 = self.convBank(x) 196 | x1 = self.convBank(x1) 197 | x2 = self.convBank(x2) 198 | 199 | # Upsample banks 1 and 2 to bank 0 size and accumulate inputs 200 | #x1 = self.upscale1(x1) 201 | #x2 = self.upscale2(x2) 202 | x1 = self.deconv1(x1) 203 | x2 = self.deconv2(x2) 204 | 205 | x = torch.cat((x0, x1, x2), dim=1) 206 | #x = x0 + x1 + x2 207 | 208 | # Apply last 2 convolutions 209 | x = F.relu(self.conv2(x)) 210 | 211 | # Output pressure (1 chan) 212 | p = self.convOut(x) 213 | 214 | 215 | # Add back the unary dimension 216 | if not self.is3D: 217 | p = torch.unsqueeze(p, 2) 218 | 219 | # Correct U = UDiv - grad(p) 220 | # flags is the one with Manta's values, not occupancy in [0,1] 221 | fluid.velocityUpdate(pressure=p, U=UDiv, flags=flags) 222 | 223 | # We now UNDO the scale factor we applied on the input. 224 | if self.mconf['normalizeInput']: 225 | p = torch.mul(p,s) # Applies p' = *= scale 226 | UDiv = torch.mul(UDiv,s) 227 | 228 | if 'periodic-x' in self.mconf and 'periodic-y' in self.mconf: 229 | U_temp = UDiv.clone() 230 | 231 | # Set BCs after velocity update. 232 | UDiv = fluid.setWallBcs(UDiv, flags) 233 | if 'periodic-x' in self.mconf and 'periodic-y' in self.mconf: 234 | if self.mconf['periodic-x']: 235 | UDiv[:,1,:,:,1] = U_temp[:,1,:,:,UDiv.size(4)-1] 236 | if self.mconf['periodic-y']: 237 | UDiv[:,0,:,1] = U_temp[:,0,:,UDiv.size(3)-1] 238 | return p, UDiv 239 | 240 | 241 | -------------------------------------------------------------------------------- /trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_conf.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_conf.pth -------------------------------------------------------------------------------- /trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_lastEpoch.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_lastEpoch.pth -------------------------------------------------------------------------------- /trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_lastEpoch_best.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_lastEpoch_best.pth -------------------------------------------------------------------------------- /trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_mconf.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/jolibrain/fluidnet_cxx/ce6ff7475d5ac16b7db23594240aa3bcd3d8c81c/trained_models/ScaleNet_ShortTerm_LongTermLoss/convModel_mconf.pth --------------------------------------------------------------------------------