adding gifs

Author: pmocz

adirondax/simulation.py

@@ -399,7 +399,7 @@ def step_fn_stacked(carry, _):
                         plot_sim(state, checkpoint_dir, i, self.params)
                 else:
                     carry, _ = jax.lax.scan(step_fn_stacked, carry, xs=None, length=nt)
-                return carry
+            return carry
 
         # save initial state
         if jax.process_index() == 0:

examples/kelvin_helmholtz/output.gif

@@ -20,7 +20,7 @@
 """
 
 
-def set_up_simulation():
+def set_up_simulation(save=False):
     # Define the parameters for the simulation
     n = 128
     nt = 100 * int(n / 128)
@@ -40,6 +40,10 @@ def set_up_simulation():
             "span": t_stop,
             "num_timesteps": nt,
         },
+        "output": {
+            "save": save,
+            "plot_dynamic_range": 2.0,
+        },
     }
 
     # Initialize the simulation
@@ -51,7 +55,7 @@ def set_up_simulation():
 def solve_inverse_problem(sim):
     # Load the target density field
     target_data = img.imread("target.png")[:, :, 0]
-    rho_target = jnp.flipud(jnp.array(target_data, dtype=float))
+    rho_target = jnp.flipud(jnp.array(target_data, dtype=float)).T
     rho_target = 1.0 - 0.5 * (rho_target - 0.5)
     rho_target /= jnp.mean(rho_target)
 
@@ -100,7 +104,7 @@ def make_plot(psi, theta):
     ax2 = plt.subplot(grid[0, 1])
     plt.sca(ax1)
     plt.cla()
-    plt.imshow(theta, cmap="bwr")
+    plt.imshow(theta.T, cmap="bwr")
     plt.clim(-jnp.pi, jnp.pi)
     ax1.get_xaxis().set_visible(False)
     ax1.get_yaxis().set_visible(False)
@@ -109,7 +113,7 @@ def make_plot(psi, theta):
     plt.title(r"${\rm initial\,angle}(\psi)$")
     plt.sca(ax2)
     plt.cla()
-    plt.imshow(jnp.log10(jnp.abs(psi) ** 2), cmap="inferno")
+    plt.imshow(jnp.log10(jnp.abs(psi) ** 2).T, cmap="inferno")
     plt.clim(-0.2, 0.2)
     ax2.get_xaxis().set_visible(False)
     ax2.get_yaxis().set_visible(False)
@@ -122,8 +126,9 @@ def make_plot(psi, theta):
 
 
 def main():
-    sim = set_up_simulation()
+    sim = set_up_simulation(save=False)
     theta = solve_inverse_problem(sim)
+    sim = set_up_simulation(save=True)
     psi = rerun_simulation(sim, theta)
     make_plot(psi, theta)
 

examples/logo_inverse_problem/logo_inverse_problem.py

@@ -102,7 +102,7 @@ def continuing_criterion(carry):
 def solve_inverse_problem(sim):
     # Load the target density field
     target_data = img.imread("target.png")[:, :, 0]
-    rho_target = jnp.flipud(jnp.array(target_data, dtype=float))
+    rho_target = jnp.flipud(jnp.array(target_data, dtype=float)).T
     rho_target = 1.0 - 0.5 * (rho_target - 0.5)
     rho_target /= jnp.mean(rho_target)
 
@@ -151,7 +151,7 @@ def make_plot(psi, theta):
     ax2 = plt.subplot(grid[0, 1])
     plt.sca(ax1)
     plt.cla()
-    plt.imshow(theta, cmap="bwr")
+    plt.imshow(theta.T, cmap="bwr")
     plt.clim(-jnp.pi, jnp.pi)
     ax1.get_xaxis().set_visible(False)
     ax1.get_yaxis().set_visible(False)
@@ -160,7 +160,7 @@ def make_plot(psi, theta):
     plt.title(r"${\rm initial\,angle}(\psi)$")
     plt.sca(ax2)
     plt.cla()
-    plt.imshow(jnp.log10(jnp.abs(psi) ** 2), cmap="inferno")
+    plt.imshow(jnp.log10(jnp.abs(psi) ** 2).T, cmap="inferno")
     plt.clim(-0.2, 0.2)
     ax2.get_xaxis().set_visible(False)
     ax2.get_yaxis().set_visible(False)

examples/logo_inverse_problem/logo_inverse_problem_optax.py

@@ -23,6 +23,7 @@
 def set_up_simulation():
     # Define the parameters for the simulation
     n = 512
+    nt = -1  # 2400
     t_stop = 0.5
     gamma = 5.0 / 3.0
     box_size = 1.0
@@ -40,6 +41,12 @@ def set_up_simulation():
         },
         "time": {
             "span": t_stop,
+            "num_timesteps": nt,
+        },
+        "output": {
+            "num_checkpoints": 100,
+            "save": False,
+            "plot_dynamic_range": 2.3,
         },
         "hydro": {
             "eos": {"type": "ideal", "gamma": gamma},

examples/logo_inverse_problem/output.gif

@@ -10,5 +10,5 @@ if [ -z "$folder" ]; then
   exit 1
 fi
 
-magick "$folder"/dm*.png "$output"
+magick "$folder"/*.png "$output"
 echo "GIF created: $output"