Improve plot_force and fix spec_updater

tests/draw/test_plt.py

@@ -15,7 +15,7 @@
 from tyssue import Sheet, config
 from tyssue.draw.plt_draw import quick_edge_draw, sheet_view
 from tyssue.draw.plt_draw import _face_color_from_sequence
-from tyssue.draw.plt_draw import create_gif
+from tyssue.draw.plt_draw import create_gif, plot_forces
 
 
 class TestsPlt:
@@ -51,7 +51,7 @@ def test_sheet_view(self):
             0.0, 1.0, num=self.sheet.edge_df.shape[0]
         )[::-1]
 
-        self.draw_specs["edge"]["visible"] = True
+        self.draw_specs["edge"]["visible"] = True 
         self.draw_specs["edge"]["color"] = self.sheet.edge_df["rand"]  # [0, 0, 0, 1]
         self.draw_specs["edge"]["alpha"] = 1.0
         self.draw_specs["edge"]["color_range"] = 0, 3
@@ -155,3 +155,19 @@ def test_create_gif():
 
     os.remove("frames.gif")
     os.remove("interval.gif")
+
+
+def test_plot_forces():
+    geom = SheetGeometry
+    model = PlanarModel
+    sheet = Sheet("3", *three_faces_sheet())
+    sheet.update_specs(model.specs) 
+    geom.update_all(sheet)
+    fig, ax = plot_forces(sheet,
+                         geom,
+                         model,
+                         list('xy'),
+                         0.05,
+                         **{'extract': {'x_boundary': (-10, 10)}})
+
+    assert ax.lines[0].get_xydata().shape == (54, 2)

tests/utils/test_utils.py

@@ -39,6 +39,11 @@ def test_spec_updater():
     print(specs)
     assert specs["face"]["geometry"] == new_specs["face"]["geometry"]
 
+    new_specs = {"extract":{"z_boundary":(-10,10)}}
+    utils.spec_updater(specs, new_specs)
+    assert "extract" in specs
+    assert specs["extract"]['z_boundary'] == new_specs["extract"]['z_boundary']
+
 
 def test_set_data_columns():
     dsets, _ = three_faces_sheet()

tyssue/draw/plt_draw.py

@@ -402,15 +402,17 @@ def plot_forces(
     else:
         grad_i = model.compute_gradient(sheet, components=False) * scaling
         grad_i = grad_i.loc[sheet.vert_df["is_active"].astype(bool)]
-    arrows = pd.DataFrame(columns=coords + gcoords, index=sheet.vert_df.index)
-    arrows[coords] = sheet.vert_df[coords]
-    arrows[gcoords] = -grad_i[gcoords]  # F = -grad E
+    sheet.vert_df[gcoords]=-grad_i[gcoords] # F = -grad E
+
+    if 'extract' in draw_specs:
+        sheet = sheet.extract_bounding_box(**draw_specs['extract'])
 
     if ax is None:
         fig, ax = quick_edge_draw(sheet, coords)
     else:
         fig = ax.get_figure()
 
+    arrows = sheet.vert_df[coords+gcoords]
     for _, arrow in arrows.iterrows():
         ax.arrow(*arrow, **draw_specs["grad"])
     return fig, ax

tyssue/utils/utils.py

@@ -58,9 +58,12 @@ def spec_updater(specs, new):
     specs: specification that will be modified
     new: dictionary of new specification
     """
-    for key, spec in specs.items():
-        if new.get(key) is not None:
-            spec.update(new[key])
+    for key in new.keys():
+        if specs.get(key) is not None:
+            specs[key].update(new[key])
+        else:
+            specs[key] = new[key]
+
 
 
 def set_data_columns(datasets, specs, reset=False):