Adding vorticity to model parameter using FiLM (Feature wise Linear Modulation)

config.yaml

@@ -14,6 +14,10 @@ simulation:
   fps: 24.0
   grid_resolution: 128
   input_channels: 6      #[density, velx, vely, density-1, emitter_mask, collider_mask]
+  vorticity_default: 0.2
+  vorticity_min: 0.1
+  vorticity_max: 0.4
+  vorticity_step: 0.1
 
 engine:
   window_width: 1250

engine/include/Config.hpp

@@ -20,6 +20,10 @@ struct SimulationConfig
     float fps{30.0f};
     int gridResolution{128};
     int inputChannels{0};
+    float vorticityDefault{0.2f};
+    float vorticityMin{0.1f};
+    float vorticityMax{0.4f};
+    float vorticityStep{0.1f};
 };
 
 struct ModelsConfig
@@ -65,6 +69,22 @@ class Config
     {
         return m_simulationConfig.inputChannels;
     }
+    float getVorticityDefault() const
+    {
+        return m_simulationConfig.vorticityDefault;
+    }
+    float getVorticityMin() const
+    {
+        return m_simulationConfig.vorticityMin;
+    }
+    float getVorticityMax() const
+    {
+        return m_simulationConfig.vorticityMax;
+    }
+    float getVorticityStep() const
+    {
+        return m_simulationConfig.vorticityStep;
+    }
 
     const std::filesystem::path& getModelsFolder() const
     {

engine/include/FluidScene.hpp

@@ -72,6 +72,11 @@ class FluidScene final : public Scene
     float m_velocityStrength{0.5f};
     float m_velocityDecayPercent{5.0f};
 
+    float m_vorticity{0.2f};
+    float m_vorticityMin{0.1f};
+    float m_vorticityMax{0.4f};
+    float m_vorticityStep{0.1f};
+
     int m_prevMouseGridX{-1};
     int m_prevMouseGridY{-1};
     bool m_mousePressed{false};

engine/include/Simulation.hpp

@@ -28,6 +28,9 @@ class Simulation
         return !m_useGpu;
     }
 
+    void setVorticity(float v);
+    float getVorticity() const;
+
     const SimulationBuffer* getLatestState() const;
     float getAvgComputeTimeMs() const;
 
@@ -58,6 +61,8 @@ class Simulation
 
     const std::atomic<SceneMaskSnapshot*>* m_sceneSnapshotPtr{nullptr};
 
+    std::atomic<float> m_vorticity{0.2f};
+
     std::atomic<float> m_avgComputeTimeMs{0.0f};
     float m_sumComputeTimeMs{0.0f};
     int m_computeTimeSamples{0};

engine/src/Config.cpp

@@ -48,6 +48,22 @@ template <> struct convert<FluidNet::SimulationConfig>
         {
             config.inputChannels = node["input_channels"].as<int>();
         }
+        if (node["vorticity_default"])
+        {
+            config.vorticityDefault = node["vorticity_default"].as<float>();
+        }
+        if (node["vorticity_min"])
+        {
+            config.vorticityMin = node["vorticity_min"].as<float>();
+        }
+        if (node["vorticity_max"])
+        {
+            config.vorticityMax = node["vorticity_max"].as<float>();
+        }
+        if (node["vorticity_step"])
+        {
+            config.vorticityStep = node["vorticity_step"].as<float>();
+        }
         return true;
     }
 };

engine/src/FluidScene.cpp

@@ -4,6 +4,7 @@
 #include "Profiling.hpp"
 #include "SceneState.hpp"
 #include <GLFW/glfw3.h>
+#include <cmath>
 #include <imgui.h>
 #include <iostream>
 
@@ -43,6 +44,10 @@ void FluidScene::onInit()
     m_renderer->initialize();
 
     m_simulationFPS = config.getSimulationFPS();
+    m_vorticity = config.getVorticityDefault();
+    m_vorticityMin = config.getVorticityMin();
+    m_vorticityMax = config.getVorticityMax();
+    m_vorticityStep = config.getVorticityStep();
 
     std::cout << "FluidScene initialized" << std::endl;
 }
@@ -128,6 +133,15 @@ void FluidScene::onRenderUI()
     ImGui::SliderInt("Velocity Brush", &m_velocityBrushSize, 1, 15);
     ImGui::SliderFloat("Velocity Strength", &m_velocityStrength, 0.1f, 1.0f, "%.1f");
     ImGui::SliderFloat("Velocity Decay %", &m_velocityDecayPercent, 1.0f, 10.0f, "%.1f%%");
+
+    ImGui::Separator();
+    if (ImGui::SliderFloat("Vorticity", &m_vorticity, m_vorticityMin, m_vorticityMax, "%.2f"))
+    {
+        m_vorticity = std::round(m_vorticity / m_vorticityStep) * m_vorticityStep;
+        if (m_simulation)
+            m_simulation->setVorticity(m_vorticity);
+    }
+
     ImGui::Checkbox("Debug Overlay (O)", &m_showDebugOverlay);
     if (ImGui::IsItemEdited() && m_renderer)
     {

engine/src/Simulation.cpp

@@ -15,6 +15,7 @@ Simulation::Simulation()
     const auto& config = Config::getInstance();
     m_targetStepTime = 1.0f / config.getSimulationFPS();
     m_useGpu = config.isGpuEnabled();
+    m_vorticity.store(config.getVorticityDefault(), std::memory_order_relaxed);
 
     int resolution = config.getGridResolution();
 
@@ -156,6 +157,16 @@ void Simulation::setSceneSnapshot(const std::atomic<SceneMaskSnapshot*>* snapsho
     m_sceneSnapshotPtr = snapshot;
 }
 
+void Simulation::setVorticity(float v)
+{
+    m_vorticity.store(v, std::memory_order_release);
+}
+
+float Simulation::getVorticity() const
+{
+    return m_vorticity.load(std::memory_order_acquire);
+}
+
 void Simulation::workerLoop_()
 {
     PROFILE_SET_THREAD_NAME("Simulation Thread");
@@ -292,20 +303,40 @@ float Simulation::runInferenceStep_(SimulationBuffer* frontBuf, SimulationBuffer
         }
 
         Ort::AllocatorWithDefaultOptions allocator;
-        auto inputName = m_ortSession->GetInputNameAllocated(0, allocator);
+        auto inputName0 = m_ortSession->GetInputNameAllocated(0, allocator);
         auto outputName = m_ortSession->GetOutputNameAllocated(0, allocator);
-
-        const char* inputNames[] = {inputName.get()};
         const char* outputNames[] = {outputName.get()};
 
+        const size_t numModelInputs = m_ortSession->GetInputCount();
+
         auto inferenceStart = Clock::now();
 
         std::vector<Ort::Value> outputTensors;
         {
             PROFILE_SCOPE_NAMED("ONNX Runtime Execute");
             PROFILE_ZONE_TEXT("Model Inference", 15);
-            outputTensors = m_ortSession->Run(Ort::RunOptions{nullptr}, inputNames, &inputTensor, 1,
-                                              outputNames, 1);
+            if (numModelInputs >= 2)
+            {
+                auto inputName1 = m_ortSession->GetInputNameAllocated(1, allocator);
+                float vorticity = m_vorticity.load(std::memory_order_acquire);
+                const int64_t condShape[] = {1};
+                auto condMemInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
+                Ort::Value condTensor =
+                    Ort::Value::CreateTensor<float>(condMemInfo, &vorticity, 1, condShape, 1);
+
+                const char* inputNames[] = {inputName0.get(), inputName1.get()};
+                std::vector<Ort::Value> inputVec;
+                inputVec.push_back(std::move(inputTensor));
+                inputVec.push_back(std::move(condTensor));
+                outputTensors = m_ortSession->Run(Ort::RunOptions{nullptr}, inputNames,
+                                                  inputVec.data(), 2, outputNames, 1);
+            }
+            else
+            {
+                const char* inputNames[] = {inputName0.get()};
+                outputTensors = m_ortSession->Run(Ort::RunOptions{nullptr}, inputNames,
+                                                  &inputTensor, 1, outputNames, 1);
+            }
         }
 
         auto inferenceEnd = Clock::now();

ml/config/model_architectures/UNet_medium.yaml

@@ -2,7 +2,7 @@ in_channels: 6
 out_channels: 3
 base_channels: 32
 depth: 3
-norm: instance
+norm: instance #group
 act: gelu
 group_norm_groups: 8
 dropout: 0.1
@@ -12,3 +12,5 @@ padding_mode: replicate  # "zeros", "reflect", "replicate", "circular"
 use_residual: true
 bottleneck_blocks: 1
 output_activation: linear_clamp
+use_film: true
+film_cond_dim: 128

ml/config/training_config.py

@@ -68,6 +68,8 @@ class TrainingConfig(BaseModel):
     use_residual: bool = True
     bottleneck_blocks: int = 1
     output_activation: OutputActivationType = "linear_clamp"
+    use_film: bool = False
+    film_cond_dim: int = 128
 
     # MLFlow settings
     mlflow_tracking_uri: str = "./mlruns"  # mlflow server

ml/dataset/npz_sequence.py

@@ -1,3 +1,4 @@
+import json
 import time
 from pathlib import Path
 
@@ -253,7 +254,14 @@ def __init__(
             raise FileNotFoundError(f"No seq_*.npz files found in {npz_dir_path}")
 
         self.num_real_sequences = len(self.seq_paths)
-        # self.num_fake_sequences = _calculate_fake_count(self.num_real_sequences, fake_empty_pct)
+
+        self._seq_scalars: list[float | None] = []
+        for path in self.seq_paths:
+            meta_path = path.with_name(path.stem + ".meta.json")
+            if meta_path.exists():
+                self._seq_scalars.append(json.loads(meta_path.read_text()).get("vorticity"))
+            else:
+                self._seq_scalars.append(None)
 
         # Load global normalization scales
         self._norm_scales: dict[str, float] | None = None
@@ -437,14 +445,16 @@ def __len__(self) -> int:
 
     def __getitem__(
         self, idx: int
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor]:
+    ) -> (
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+    ):
         si, t = self._index[idx]
 
-        # Handle real sequences
         path = self.seq_paths[si]
+        cond_val = self._seq_scalars[si]
+        cond = torch.tensor(cond_val if cond_val is not None else 0.0, dtype=torch.float32)
 
         if self.rollout_steps > 1:
-            # Multi-step rollout mode
             if self.preload and self._preloaded_sequences is not None:
                 x, y_seq, masks = self._load_rollout_sample_from_memory(si, t)
             else:
@@ -455,9 +465,8 @@ def __getitem__(
 
                 x, y_seq, masks = apply_rollout_augmentation(x, y_seq, masks, self.flip_probability)
 
-            return x, y_seq, masks
+            return x, y_seq, masks, cond
         else:
-            # Single-step mode (existing code)
             if self.preload and self._preloaded_sequences is not None:
                 x, y = self._load_sample_from_memory(si, t)
             else:
@@ -466,4 +475,4 @@ def __getitem__(
             if self.is_training and self.enable_augmentation:
                 x, y = apply_augmentation(x, y, self.flip_probability)
 
-            return x, y
+            return x, y, cond