Add `block` argument to hooke.plugin.vclamp.scale and only run required commands

diff --git a/hooke/plugin/vclamp.py b/hooke/plugin/vclamp.py
index 2365122aaff8a278397e62439ff845cda15e1fb0..9151765b296bd2e5d621d1dc48bf095370618875 100644 (file)
--- a/hooke/plugin/vclamp.py
+++ b/hooke/plugin/vclamp.py
@@ -38,22 +38,36 @@ from ..util.fit import PoorFit, ModelFitter
 from .curve import CurveArgument
 
 
-def scale(hooke, curve):
+def scale(hooke, curve, block=None):
+    """Run 'add block force array' on `block`.
+
+    Runs 'zero block surface contact point' first, if necessary.  Does
+    not run either command if the columns they add to the block are
+    already present.
+
+    If `block` is `None`, scale all blocks in `curve`.
+    """
     commands = hooke.commands
     contact = [c for c in hooke.commands
                if c.name == 'zero block surface contact point'][0]
     force = [c for c in hooke.commands if c.name == 'add block force array'][0]
     inqueue = None
     outqueue = NullQueue()
-    for i,block in enumerate(curve.data):
-        numpy.savetxt(open('curve.dat', 'w'), block, delimiter='\t')
-        params = {'curve':curve, 'block':i}
-        try:
-            contact._run(hooke, inqueue, outqueue, params)
-        except PoorFit, e:
-            raise PoorFit('Could not fit %s %s: %s'
-                          % (curve.path, i, str(e)))
-        force._run(hooke, inqueue, outqueue, params)
+    if block == None:
+        for i in range(len(curve.data)):
+            scale(hooke, curve, block=i)
+    else:
+        params = {'curve':curve, 'block':block}
+        b = curve.data[block]
+        if ('surface distance (m)' not in b.info.columns
+            or 'surface adjusted deflection (m)' not in b.info.columns):
+            try:
+                contact._run(hooke, inqueue, outqueue, params)
+            except PoorFit, e:
+                raise PoorFit('Could not fit %s %s: %s'
+                              % (curve.path, block, str(e)))
+        if ('deflection (N)' not in b.info.columns):
+            force._run(hooke, inqueue, outqueue, params)
     return curve
 
 class SurfacePositionModel (ModelFitter):