Update Shallow Water PDE example to use GIF writer (#684)

* Render shallow water PDE example visualizations with GIF writer

- Use Tensor as protocol requirement for water level output
- Remove single-purpose visualization code
- Updated readme file

* GIF quantizing for grayscale image tensors

- Greyscale image is represented as [width, height, 1] shaped tensor
- New palette with 64 shades of grey

* Improve CLI of shallow water PDE example

- Options for setting resolution, duration, iterations, learning rate and target image
- Write animations to 'output' directory
- Can be run from repository root
- Updated readme

Author: vojtamolda

Examples/Shallow-Water-PDE/ArrayLoopSolution.swift

@@ -43,7 +43,9 @@ import TensorFlow
 ///
 struct ArrayLoopSolution: ShallowWaterEquationSolution {
   /// Water level height
-  var waterLevel: [[Float]] { u1 }
+  var waterLevel: Tensor<Float> {
+    Tensor<Float>(shape: TensorShape([resolution, resolution]), scalars: u1.flatMap{ $0 })
+  }
   /// Solution time
   var time: Float { t }
 
@@ -58,7 +60,7 @@ struct ArrayLoopSolution: ShallowWaterEquationSolution {
   /// Dispersion coefficient
   @noDerivative private let α: Float = 0.00001
   /// Number of spatial grid points
-  @noDerivative private let resolution: Int = 256
+  @noDerivative private let resolution: Int
   /// Spatial discretization step
   @noDerivative private var Δx: Float { 1 / Float(resolution) }
   /// Time-step calculated to stay below the CFL stability limit
@@ -67,11 +69,11 @@ struct ArrayLoopSolution: ShallowWaterEquationSolution {
   /// Creates initial solution with water level `u0` at time `t`.
   @differentiable
   init(waterLevel u0: [[Float]], time t: Float = 0.0) {
+    self.resolution = u0.count
     self.u0 = u0
     self.u1 = u0
     self.t = t
 
-    precondition(u0.count == resolution)
     precondition(u0.allSatisfy { $0.count == resolution })
   }
 
@@ -102,11 +104,11 @@ struct ArrayLoopSolution: ShallowWaterEquationSolution {
   /// Constructs intermediate solution with previous water level `u0`, current water level `u1` and time `t`.
   @differentiable
   private init(u0: [[Float]], u1: [[Float]], t: Float) {
+    self.resolution = u0.count
     self.u0 = u0
     self.u1 = u1
     self.t = t
 
-    precondition(u0.count == self.resolution)
     precondition(u0.allSatisfy { $0.count == self.resolution })
     precondition(u1.count == self.resolution)
     precondition(u1.allSatisfy { $0.count == self.resolution })

Examples/Shallow-Water-PDE/Images/Optimization.gif

@@ -9,43 +9,33 @@ The code is re-implemented from the [DiffTaichi paper](https://arxiv.org/abs/191
 
 ## Splash in a Bathtub
 
-The following code builds and runs a demo simulation of a water surface behavior in a rectangular bathtub. There's an initial "splash" at the beginning that drives the simulation. The splash generates surface gravity waves that propagate away from the center and reflect off the domain walls.
-
 <img alt="Simulation of a splash in bathtub" src="Images/Splash.gif" align="right">
 
-```sh
-# Make sure you run the example from this directory
-cd swift-models/Examples/Shallow-Water-PDE
+The following code builds and runs a demo simulation of a water surface behavior in a rectangular bathtub. There's an initial "splash" at the beginning that drives the simulation. The splash generates surface gravity waves that propagate away from the center and reflect off the domain walls.
 
+```sh
 # Build and run the example app with the splash flag
 swift run -c release Shallow-Water-PDE --splash
-
-# Create animation and remove the intermediate files
-convert Images/Splash-*.jpg Images/Splash.gif
-rm Images/Splash-*.jpg
 ```
 
+Animation of the solution is saved to ` output/splash.gif`.
+
 
 ## Optimization of Initial Water Level
 
+<img alt="Result of initial condition optimization" src="Images/Optimization.gif" align="right">
+
 The following example takes advantage of the end-to-end differentiability of the PDE solution and optimizes the initial condition. The initial condition is the water surface height used to start the simulation. The goal, a very tricky one from the perspective of optimal control theory, is to achieve specific wave pattern when the simulation terminates.
 
 The goal is represented by a MSE cost function that measures difference between the terminal water level and a target picture. Optimization procedure itself is a simple gradient descent that adjust the initial water surface height to achieve smaller cost after every iteration.
 
-<img alt="Result of initial condition optimization" src="Images/Optimization.gif" align="right">
-
 ```sh
-# Make sure you run the example from this directory
-cd swift-models/Examples/Shallow-Water-PDE
-
 # Build and run the example app with the optimization flag
 swift run -c release Shallow-Water-PDE --optimization
-
-# Create animation and remove the intermediate files
-convert Images/Optimization-*.jpg Images/Optimization.gif
-rm Images/Optimization-*.jpg
 ```
 
+Animation of the optimized solution is saved to `output/optimization.gif`.
+
 
 ## Benchmark of 4 Solver Implementations
 

Examples/Shallow-Water-PDE/Images/Splash.gif

@@ -19,7 +19,7 @@ import TensorFlow
 /// Differentiable solution of shallow water equation on a unit square.
 protocol ShallowWaterEquationSolution: Differentiable {
   /// Snapshot of water level height at time `time`.
-  @noDerivative var waterLevel: [[Float]] { get }
+  @noDerivative var waterLevel: Tensor<Float> { get }
   /// Solution time
   @noDerivative var time: Float { get }
 
@@ -44,4 +44,26 @@ extension Array where Array.Element: ShallowWaterEquationSolution {
     }
     self.append(currentSolution)
   }
+
+  /// Saves the shallow water equation solutions as a GIF image with the specified `delay` and keeping every `keep` frames.
+  func saveAnimatedImage(directory: String, name: String, delay: Int = 4, keep: Int = 8) throws {
+    let filtered = self.enumerated().filter { $0.offset % keep == 0 }.map { $0.element }
+    let frames = filtered.map { $0.waterLevel.normalizedGrayscaleImage(min: -1, max: +1) }
+    try frames.saveAnimatedImage(directory: directory, name: name, delay: delay)
+  }
+}
+
+// MARK: - Utilities
+
+extension Tensor where Scalar == Float {
+  /// Returns a 3D grayscale image tensor clipped and normalized from `min`-`max` to 0-255  range.
+  func normalizedGrayscaleImage(min: Scalar = -1, max: Scalar = +1) -> Tensor<Float> {
+    precondition(max > min)
+    precondition(rank == 2)
+
+    let clipped = self.clipped(min: min, max: max)
+    let normalized = (clipped - min) / (max - min) * Float(255.0)
+    
+    return normalized.expandingShape(at: 2)
+  }
 }

Examples/Shallow-Water-PDE/README.md

@@ -49,7 +49,7 @@ import TensorFlow
 ///
 struct TensorConvSolution: ShallowWaterEquationSolution {
   /// Water level height
-  var waterLevel: [[Float]] { u1.array.map { $0.scalars } }
+  var waterLevel: Tensor<Float> { u1 }
   /// Solution time
   var time: Float { t }
 
@@ -64,7 +64,7 @@ struct TensorConvSolution: ShallowWaterEquationSolution {
   /// Dispersion coefficient
   @noDerivative private let α: Float = 0.00001
   /// Number of spatial grid points
-  @noDerivative private let resolution: Int = 256
+  @noDerivative private let resolution: Int
   /// Spatial discretization step
   @noDerivative private var Δx: Float { 1 / Float(resolution) }
   /// Time-step calculated to stay below the CFL stability limit
@@ -76,6 +76,7 @@ struct TensorConvSolution: ShallowWaterEquationSolution {
   /// Creates initial solution with water level `u0` at time `t`.
   @differentiable
   init(waterLevel u0: Tensor<Float>, time t: Float = 0.0) {
+    self.resolution = u0.shape[0]
     self.u0 = u0
     self.u1 = u0
     self.t = t
@@ -125,6 +126,7 @@ struct TensorConvSolution: ShallowWaterEquationSolution {
   /// Constructs intermediate solution with previous water level `u0`, current water level `u1` and time `t`.
   @differentiable
   private init(u0: Tensor<Float>, u1: Tensor<Float>, t: Float) {
+    self.resolution = u0.shape[0]
     self.u0 = u0
     self.u1 = u1
     self.t = t

Examples/Shallow-Water-PDE/Solution.swift

@@ -44,7 +44,7 @@ import TensorFlow
 ///
 struct TensorLoopSolution: ShallowWaterEquationSolution {
   /// Water level height
-  var waterLevel: [[Float]] { u1.array.map { $0.scalars } }
+  var waterLevel: Tensor<Float> { u1 }
   /// Solution time
   var time: Float { t }
 
@@ -59,7 +59,7 @@ struct TensorLoopSolution: ShallowWaterEquationSolution {
   /// Dispersion coefficient
   @noDerivative private let α: Float = 0.00001
   /// Number of spatial grid points
-  @noDerivative private let resolution: Int = 256
+  @noDerivative private let resolution: Int
   /// Spatial discretization step
   @noDerivative private var Δx: Float { 1 / Float(resolution) }
   /// Time-step calculated to stay below the CFL stability limit
@@ -68,6 +68,7 @@ struct TensorLoopSolution: ShallowWaterEquationSolution {
   /// Creates initial solution with water level `u0` at time `t` using the specified TensorFlow `device`.
   @differentiable
   init(waterLevel u0: Tensor<Float>, time t: Float = 0.0) {
+    self.resolution = u0.shape[0]
     self.u0 = u0
     self.u1 = u0
     self.t = t
@@ -112,6 +113,7 @@ struct TensorLoopSolution: ShallowWaterEquationSolution {
   /// Constructs intermediate solution with previous water level `u0`, current water level `u1` and time `t`.
   @differentiable
   private init(u0: Tensor<Float>, u1: Tensor<Float>, t: Float) {
+    self.resolution = u0.shape[0]
     self.u0 = u0
     self.u1 = u1
     self.t = t

Examples/Shallow-Water-PDE/TensorConvSolution.swift

@@ -49,7 +49,7 @@ import TensorFlow
 ///
 struct TensorSliceSolution: ShallowWaterEquationSolution {
   /// Water level height
-  var waterLevel: [[Float]] { u1.array.map { $0.scalars } }
+  var waterLevel: Tensor<Float> { u1 }
   /// Solution time
   var time: Float { t }
 
@@ -64,7 +64,7 @@ struct TensorSliceSolution: ShallowWaterEquationSolution {
   /// Dispersion coefficient
   @noDerivative private let α: Float = 0.00001
   /// Number of spatial grid points
-  @noDerivative private let resolution: Int = 256
+  @noDerivative private let resolution: Int
   /// Spatial discretization step
   @noDerivative private var Δx: Float { 1 / Float(resolution) }
   /// Time-step calculated to stay below the CFL stability limit
@@ -73,6 +73,7 @@ struct TensorSliceSolution: ShallowWaterEquationSolution {
   /// Creates initial solution with water level `u0` at time `t`.
   @differentiable
   init(waterLevel u0: Tensor<Float>, time t: Float = 0.0) {
+    self.resolution = u0.shape[0]
     self.u0 = u0
     self.u1 = u0
     self.t = t
@@ -116,6 +117,7 @@ struct TensorSliceSolution: ShallowWaterEquationSolution {
   /// Constructs intermediate solution with previous water level `u0`, current water level `u1` and time `t`.
   @differentiable
   private init(u0: Tensor<Float>, u1: Tensor<Float>, t: Float) {
+    self.resolution = u0.shape[0]
     self.u0 = u0
     self.u1 = u1
     self.t = t