Update tools.py

tools.py

@@ -4,7 +4,7 @@ import numpy as np
 
 def plot_3D_invfitness(trait, fitness, resident, range, color="RdBu"):
     X, Y = np.meshgrid(trait, trait)
-    f_projection = (np.min(fitness) - 2) * np.ones(fitness.shape)
+    f_projection = (np.min(fitness) - np.mean(fitness)) * np.ones(fitness.shape)
     axis = dict(
         showbackground=True,
         backgroundcolor="rgb(230, 230,230)",
@@ -15,15 +15,15 @@ def plot_3D_invfitness(trait, fitness, resident, range, color="RdBu"):
 
     layout = go.Layout(
         autosize=False,
-        width=600,
-        height=500,
+        width=700,
+        height=600,
         scene=dict(
             xaxis=dict(axis),
             yaxis=dict(axis),
             zaxis=dict(axis, range=range),
             aspectratio=dict(x=1, y=1, z=1),
-            xaxis_title="Resident trait",
-            yaxis_title="Mutant trait",
+            xaxis_title="Resident trait (z_r)",
+            yaxis_title="Mutant trait (z_m)",
             zaxis_title="Invasion fitness",
         ),
     )
@@ -44,7 +44,7 @@ def plot_3D_invfitness(trait, fitness, resident, range, color="RdBu"):
         z=(fitness * 1e5),
         surfacecolor=x_projection,
         colorscale="Greys",
-        opacity=0.5,
+        opacity=0.7,
         showlegend=False,
         showscale=False,
     )
@@ -63,29 +63,28 @@ def plot_invasionfitness(zm, zlist, fitness_func, pars, range):
     inv_fitness = fitness_func(zm, zlist, pars)
 
     fig = px.line(
-        x=zlist, y=inv_fitness, labels={
-            "x": "Mutant trait value (z)", "y": "Invasion fitness"}
+        x=zlist, y=inv_fitness, labels={"x": "Mutant trait value (z_m)", "y": "Invasion fitness"}
     )
     fig.add_vline(x=zm, line_dash="dashdot")
     fig.add_hline(y=0, line_dash="dash")
     fig.update_layout(
         title="Interactive invasion process",
         xaxis=dict(range=[0, zlist[-1]], autorange=False),
-        yaxis=dict(range=range, autorange=False), autosize=False,
+        yaxis=dict(range=range, autorange=False),
+        autosize=False,
         width=450,
-        height=400
+        height=400,
     )
     return fig
 
 
-def make_interactive_video(z_start, z_end, step, zlist, fitness_func, pars, range):
+def make_interactive_video(z_start, z_end, steps, zlist, fitness_func, pars, range):
     inv_vid = []
-    for z_val in np.arange(z_start, z_end, step):
+    for z_val in np.linspace(z_start, z_end, steps):
         inv_vid.append(fitness_func(z_val, zlist, pars))
     vid = go.Figure(
         data=[
-            go.Line(x=zlist, y=fitness_func(
-                z_start, zlist, pars), name="invasion fitness"),
+            go.Line(x=zlist, y=fitness_func(z_start, zlist, pars), name="invasion fitness"),
             go.Line(
                 x=zlist,
                 y=np.zeros(len(zlist)),
@@ -101,21 +100,42 @@ def make_interactive_video(z_start, z_end, step, zlist, fitness_func, pars, rang
             ),
         ],
         layout=go.Layout(
-            title="Invasion process video", autosize=False,
-            width=550,
+            autosize=False,
+            width=650,
             height=500,
             xaxis=dict(range=[0, zlist[-1]], autorange=False),
             yaxis=dict(range=range, autorange=False),
-            xaxis_title="Mutant trait value",
+            xaxis_title="Mutant trait value (z_m)",
             updatemenus=[
-                dict(type="buttons", buttons=[
-                     dict(label="Play", method="animate", args=[None, {"frame": {"duration": 500, "redraw": False},
-                                                                       "fromcurrent": True, "transition": {"duration": 300,
-                                                                                                           "easing": "quadratic-in-out"}}]),
-                     dict(label="Pause", method="animate", args=[[None], {"frame": {"duration": 0, "redraw": False},
-                                                                          "mode": "immediate",
-                                                                          "transition": {"duration": 0}}])]),
-
+                dict(
+                    type="buttons",
+                    buttons=[
+                        dict(
+                            label="Play",
+                            method="animate",
+                            args=[
+                                None,
+                                {
+                                    "frame": {"duration": 500, "redraw": False},
+                                    "fromcurrent": True,
+                                    "transition": {"duration": 300, "easing": "quadratic-in-out"},
+                                },
+                            ],
+                        ),
+                        dict(
+                            label="Pause",
+                            method="animate",
+                            args=[
+                                [None],
+                                {
+                                    "frame": {"duration": 0, "redraw": False},
+                                    "mode": "immediate",
+                                    "transition": {"duration": 0},
+                                },
+                            ],
+                        ),
+                    ],
+                ),
             ],
         ),
         frames=[
@@ -135,7 +155,7 @@ def make_interactive_video(z_start, z_end, step, zlist, fitness_func, pars, rang
                     ),
                 ]
             )
-            for i, z_val in zip(inv_vid, np.arange(z_start, z_end, step))
+            for i, z_val in zip(inv_vid, np.linspace(z_start, z_end, steps))
         ],
     )
     return vid
@@ -162,7 +182,8 @@ def plot_PIP(zlist, fitness_func, pars):
         autosize=False,
         width=400,
         height=500,
-        xaxis_title="Resident trait",
-        yaxis_title="Mutant trait",
+        xaxis_title=r"Resident trait (z_r)",
+        yaxis_title=r"Mutant trait (z_m)",
+        title="Pairwise invasibility plot (PIP)",
     )
     return fig