update example cases with run-flag for documentation gallery

Author: skim0119

examples/AxialStretchingCase/README.md

@@ -1,34 +1,23 @@
-"""Axial stretching test-case
-
-Assume we have a rod lying aligned in the x-direction, with high internal
-damping.
-
-We fix one end (say, the left end) of the rod to a wall. On the right
-end we apply a force directed axially pulling the rods tip. Linear
-theory (assuming small displacements) predict that the net displacement
-experienced by the rod tip is Δx = FL/AE where the symbols carry their
-usual meaning (the rod is just a linear spring). We compare our results
-with the above result.
-
-We can "improve" the theory by having a better estimate for the rod's
-spring constant by assuming that it equilibriates under the new position,
-with
-Δx = F * (L + Δx)/ (A * E)
-which results in Δx = (F*l)/(A*E - F). Our rod reaches equilibrium wrt to
-this position.
-
-Note that if the damping is not high, the rod oscillates about the eventual
-resting position (and this agrees with the theoretical predictions without
-any damping : we should see the rod oscillating simple-harmonically in time).
-
-isort:skip_file
 """
+Axial Stretching
+================
+
+This case tests the axial stretching of a rod. A rod is fixed at one end and
+a force is applied at the other end. The rod stretches and the displacement
+of the tip is compared with the analytical solution.
+"""
+
+# isort:skip_file
 
 import numpy as np
 from matplotlib import pyplot as plt
 
 import elastica as ea
 
+# %%
+# First, we define a simulator class that inherits from the necessary mixins.
+# This makes it easy to add constraints, forces, and damping to the system.
+
 
 class StretchingBeamSimulator(
     ea.BaseSystemCollection, ea.Constraints, ea.Forcing, ea.Damping, ea.CallBacks
@@ -39,10 +28,10 @@ class StretchingBeamSimulator(
 stretch_sim = StretchingBeamSimulator()
 final_time = 200.0
 
-# Options
-PLOT_FIGURE = True
-SAVE_FIGURE = False
-SAVE_RESULTS = False
+# %%
+# Next, we set up the test parameters for the simulation. This includes the
+# number of elements, the start position, direction, normal, length, radius,
+# density, and Young's modulus of the rod.
 
 # setting up test params
 n_elem = 19
@@ -58,6 +47,10 @@ class StretchingBeamSimulator(
 poisson_ratio = 0.5
 shear_modulus = youngs_modulus / (poisson_ratio + 1.0)
 
+# %%
+# Now we can create the `CosseratRod` object. We use the `straight_rod` method
+# to create a straight rod with the specified parameters.
+
 stretchable_rod = ea.CosseratRod.straight_rod(
     n_elem,
     start,
@@ -71,16 +64,29 @@ class StretchingBeamSimulator(
 )
 
 stretch_sim.append(stretchable_rod)
+
+# %%
+# We then apply a boundary condition to fix one end of the rod. We use the
+# `OneEndFixedBC` constraint to fix the position and director of the first node.
+
 stretch_sim.constrain(stretchable_rod).using(
     ea.OneEndFixedBC, constrained_position_idx=(0,), constrained_director_idx=(0,)
 )
 
+# %%
+# A force is applied to the other end of the rod. We use the `EndpointForces`
+# forcing to apply a force in the x-direction.
+
 end_force_x = 1.0
 end_force = np.array([end_force_x, 0.0, 0.0])
 stretch_sim.add_forcing_to(stretchable_rod).using(
     ea.EndpointForces, 0.0 * end_force, end_force, ramp_up_time=1e-2
 )
 
+# %%
+# Damping is added to the system to help it reach a steady state. We use an
+# `AnalyticalLinearDamper` to add damping to the rod.
+
 # add damping
 dl = base_length / n_elem
 dt = 0.1 * dl
@@ -92,6 +98,11 @@ class StretchingBeamSimulator(
 )
 
 
+# %%
+# We define a callback class to record the position and velocity of the rod
+# during the simulation. This is useful for post-processing the results.
+
+
 # Add call backs
 class AxialStretchingCallBack(ea.CallBackBaseClass):
     """
@@ -125,54 +136,41 @@ def make_callback(
     AxialStretchingCallBack, step_skip=200, callback_params=recorded_history
 )
 
+# %%
+# We finalize the simulator and create the time-stepper. The `PositionVerlet`
+# time-stepper is used to integrate the system.
+
 stretch_sim.finalize()
 timestepper: ea.typing.StepperProtocol = ea.PositionVerlet()
 # timestepper = PEFRL()
 
+# %%
+# The simulation is run for the specified `final_time`.
+
 total_steps = int(final_time / dt)
 print("Total steps", total_steps)
 dt = final_time / total_steps
 time = 0.0
 for i in range(total_steps):
     time = timestepper.step(stretch_sim, time, dt)
 
-if PLOT_FIGURE:
-    # First-order theory with base-length
-    expected_tip_disp = end_force_x * base_length / base_area / youngs_modulus
-    # First-order theory with modified-length, gives better estimates
-    expected_tip_disp_improved = (
-        end_force_x * base_length / (base_area * youngs_modulus - end_force_x)
-    )
-
-    fig = plt.figure(figsize=(10, 8), frameon=True, dpi=150)
-    ax = fig.add_subplot(111)
-    ax.plot(recorded_history["time"], recorded_history["position"], lw=2.0)
-    ax.hlines(base_length + expected_tip_disp, 0.0, final_time, "k", "dashdot", lw=1.0)
-    ax.hlines(
-        base_length + expected_tip_disp_improved, 0.0, final_time, "k", "dashed", lw=2.0
-    )
-    if SAVE_FIGURE:
-        fig.savefig("axial_stretching.pdf")
-    plt.show()
-
-if SAVE_RESULTS:
-    import pickle
-
-    filename = "axial_stretching_data.dat"
-    file = open(filename, "wb")
-    pickle.dump(stretchable_rod, file)
-    file.close()
-
-    tv = (
-        np.asarray(recorded_history["time"]),
-        np.asarray(recorded_history["velocity_norms"]),
-    )
-
-    def as_time_series(v: np.ndarray) -> np.ndarray:
-        return v.T
-
-    np.savetxt(
-        "velocity_norms.csv",
-        as_time_series(np.stack(tv)),
-        delimiter=",",
-    )
+# %%
+# Finally, we plot the results and compare them with the analytical solution.
+# The analytical solution is calculated using the first-order theory with
+# both the base length and the modified length.
+
+# First-order theory with base-length
+expected_tip_disp = end_force_x * base_length / base_area / youngs_modulus
+# First-order theory with modified-length, gives better estimates
+expected_tip_disp_improved = (
+    end_force_x * base_length / (base_area * youngs_modulus - end_force_x)
+)
+
+fig = plt.figure(figsize=(10, 8), frameon=True, dpi=150)
+ax = fig.add_subplot(111)
+ax.plot(recorded_history["time"], recorded_history["position"], lw=2.0)
+ax.hlines(base_length + expected_tip_disp, 0.0, final_time, "k", "dashdot", lw=1.0)
+ax.hlines(
+    base_length + expected_tip_disp_improved, 0.0, final_time, "k", "dashed", lw=2.0
+)
+plt.show()

…/AxialStretchingCase/axial_stretching.py …alStretchingCase/run_axial_stretching.pyexamples/AxialStretchingCase/axial_stretching.py renamed to examples/AxialStretchingCase/run_axial_stretching.py examples/AxialStretchingCase/axial_stretching.py renamed to examples/AxialStretchingCase/run_axial_stretching.py

@@ -1,3 +1,12 @@
+"""
+Butterfly
+=========
+
+This case simulates the motion of a rod that is initially shaped like a
+butterfly. The rod is released from rest and allowed to fall under gravity.
+The simulation tracks the position and energy of the rod over time.
+"""
+
 import numpy as np
 from matplotlib import pyplot as plt
 from matplotlib.colors import to_rgb
@@ -6,6 +15,9 @@
 import elastica as ea
 from elastica.utils import MaxDimension
 
+# %%
+# First, we define a simulator class that inherits from the necessary mixins.
+
 
 class ButterflySimulator(ea.BaseSystemCollection, ea.CallBacks):
     pass
@@ -14,11 +26,10 @@ class ButterflySimulator(ea.BaseSystemCollection, ea.CallBacks):
 butterfly_sim = ButterflySimulator()
 final_time = 40.0
 
-# Options
-PLOT_FIGURE = True
-SAVE_FIGURE = True
-SAVE_RESULTS = True
-ADD_UNSHEARABLE_ROD = False
+# %%
+# Next, we set up the test parameters for the simulation. This includes the
+# number of elements, the origin, the angle of inclination, the length,
+# radius, density, and Young's modulus of the rod.
 
 # setting up test params
 # FIXME : Doesn't work with elements > 10 (the inverse rotate kernel fails)
@@ -44,6 +55,10 @@ class ButterflySimulator(ea.BaseSystemCollection, ea.CallBacks):
 poisson_ratio = 0.5
 shear_modulus = youngs_modulus / (poisson_ratio + 1.0)
 
+# %%
+# We then define the initial positions of the nodes of the rod to create the
+# butterfly shape.
+
 positions = np.empty((MaxDimension.value(), n_elem + 1))
 dl = total_length / n_elem
 
@@ -60,6 +75,9 @@ class ButterflySimulator(ea.BaseSystemCollection, ea.CallBacks):
     - np.sin(angle_of_inclination) * vertical_direction
 )
 
+# %%
+# Now we can create the `CosseratRod` object with the specified positions.
+
 butterfly_rod = ea.CosseratRod.straight_rod(
     n_elem,
     start=origin.reshape(3),
@@ -76,6 +94,11 @@ class ButterflySimulator(ea.BaseSystemCollection, ea.CallBacks):
 butterfly_sim.append(butterfly_rod)
 
 
+# %%
+# We define a callback class to record the position and energy of the rod
+# during the simulation.
+
+
 # Add call backs
 class VelocityCallBack(ea.CallBackBaseClass):
     """
@@ -117,12 +140,17 @@ def make_callback(
     VelocityCallBack, step_skip=100, callback_params=recorded_history
 )
 
+# %%
+# We finalize the simulator and create the time-stepper.
 
 butterfly_sim.finalize()
 timestepper: ea.typing.StepperProtocol
 timestepper = ea.PositionVerlet()
 # timestepper = PEFRL()
 
+# %%
+# The simulation is run for the specified `final_time`.
+
 dt = 0.01 * dl
 total_steps = int(final_time / dt)
 print("Total steps", total_steps)
@@ -131,52 +159,43 @@ def make_callback(
 for i in range(total_steps):
     time = timestepper.step(butterfly_sim, time, dt)
 
-if PLOT_FIGURE:
-    # Plot the histories
-    fig = plt.figure(figsize=(5, 4), frameon=True, dpi=150)
-    ax = fig.add_subplot(111)
-    positions_history = recorded_history["position"]
-    # record first position
-    first_position = positions_history.pop(0)
-    ax.plot(first_position[2, ...], first_position[0, ...], "r--", lw=2.0)
-    n_positions = len(positions_history)
-    for i, pos in enumerate(positions_history):
-        alpha = np.exp(i / n_positions - 1)
-        ax.plot(pos[2, ...], pos[0, ...], "b", lw=0.6, alpha=alpha)
-    # final position is also separate
-    last_position = positions_history.pop()
-    ax.plot(last_position[2, ...], last_position[0, ...], "k--", lw=2.0)
-    # don't block
-    fig.show()
-
-    # Plot the energies
-    energy_fig = plt.figure(figsize=(5, 4), frameon=True, dpi=150)
-    energy_ax = energy_fig.add_subplot(111)
-    times = np.asarray(recorded_history["time"])
-    te = np.asarray(recorded_history["te"])
-    re = np.asarray(recorded_history["re"])
-    be = np.asarray(recorded_history["be"])
-    se = np.asarray(recorded_history["se"])
-
-    energy_ax.plot(times, te, c=to_rgb("xkcd:reddish"), lw=2.0, label="Translations")
-    energy_ax.plot(times, re, c=to_rgb("xkcd:bluish"), lw=2.0, label="Rotation")
-    energy_ax.plot(times, be, c=to_rgb("xkcd:burple"), lw=2.0, label="Bend")
-    energy_ax.plot(times, se, c=to_rgb("xkcd:goldenrod"), lw=2.0, label="Shear")
-    energy_ax.plot(times, te + re + be + se, c="k", lw=2.0, label="Total energy")
-    energy_ax.legend()
-    # don't block
-    energy_fig.show()
-
-    if SAVE_FIGURE:
-        fig.savefig("butterfly.png")
-        energy_fig.savefig("energies.png")
-
-    plt.show()
-
-if SAVE_RESULTS:
-    import pickle
-
-    filename = "butterfly_data.dat"
-    file = open(filename, "wb")
-    pickle.dump(butterfly_rod, file)
-    file.close()
+# %%
+# Finally, we plot the results. The position of the rod is plotted at
+# different time steps, and the energies are plotted as a function of time.
+
+# Plot the histories
+fig = plt.figure(figsize=(5, 4), frameon=True, dpi=150)
+ax = fig.add_subplot(111)
+positions_history = recorded_history["position"]
+# record first position
+first_position = positions_history.pop(0)
+ax.plot(first_position[2, ...], first_position[0, ...], "r--", lw=2.0)
+n_positions = len(positions_history)
+for i, pos in enumerate(positions_history):
+    alpha = np.exp(i / n_positions - 1)
+    ax.plot(pos[2, ...], pos[0, ...], "b", lw=0.6, alpha=alpha)
+# final position is also separate
+last_position = positions_history.pop()
+ax.plot(last_position[2, ...], last_position[0, ...], "k--", lw=2.0)
+# don't block
+fig.show()
+
+# Plot the energies
+energy_fig = plt.figure(figsize=(5, 4), frameon=True, dpi=150)
+energy_ax = energy_fig.add_subplot(111)
+times = np.asarray(recorded_history["time"])
+te = np.asarray(recorded_history["te"])
+re = np.asarray(recorded_history["re"])
+be = np.asarray(recorded_history["be"])
+se = np.asarray(recorded_history["se"])
+
+energy_ax.plot(times, te, c=to_rgb("xkcd:reddish"), lw=2.0, label="Translations")
+energy_ax.plot(times, re, c=to_rgb("xkcd:bluish"), lw=2.0, label="Rotation")
+energy_ax.plot(times, be, c=to_rgb("xkcd:burple"), lw=2.0, label="Bend")
+energy_ax.plot(times, se, c=to_rgb("xkcd:goldenrod"), lw=2.0, label="Shear")
+energy_ax.plot(times, te + re + be + se, c="k", lw=2.0, label="Total energy")
+energy_ax.legend()
+# don't block
+energy_fig.show()
+
+plt.show()