Add 'scaled column addition' command.

[hooke.git] / hooke / plugin / curve.py

diff --git a/hooke/plugin/curve.py b/hooke/plugin/curve.py
index 821d33370e6051657cfdb298403e344d86430e5e..f6e5892691df4227dcfdf6cea8740c780020158b 100644 (file)
--- a/hooke/plugin/curve.py
+++ b/hooke/plugin/curve.py
@@ -210,6 +210,8 @@ class ColumnAccessCommand (BlockCommand):
         if column_name == None:
             column_name = self._column_arguments[0].name
         column_name = params[column_name]
+        if column_name is None:
+            return None
         block = self._block(hooke, params, block_name)
         columns = block.info['columns']
         try:
@@ -269,7 +271,7 @@ class CurvePlugin (Builtin):
             GetCommand(self), InfoCommand(self), BlockInfoCommand(self),
             DeltaCommand(self), ExportCommand(self), DifferenceCommand(self),
             DerivativeCommand(self), PowerSpectrumCommand(self),
-            ClearStackCommand(self)]
+            ScaledColumnAdditionCommand(self), ClearStackCommand(self)]
 
 
 # Define commands
@@ -686,6 +688,92 @@ Otherwise, the chunks are end-to-end, and not overlapping.
         return params
 
 
+class ScaledColumnAdditionCommand (ColumnAddingCommand):
+    """Add one affine transformed column to another: `o=A*i1+B*i2+C`.
+    """
+    def __init__(self, plugin):
+        super(ScaledColumnAdditionCommand, self).__init__(
+            name='scaled column addition',
+            columns=[
+                ('input column 1', 'input column (m)', """
+Name of the first column to use as the transform input.
+""".strip()),
+                ('input column 2', None, """
+Name of the second column to use as the transform input.
+""".strip()),
+                ],
+            new_columns=[
+                ('output column', 'output column (m)', """
+Name of the column to use as the transform output.
+""".strip()),
+                ],
+            arguments=[
+                Argument(name='scale 1', type='float', default=None,
+                         help="""
+A float value for the first scale constant.
+""".strip()),
+                Argument(name='scale 1 name', type='string', default=None,
+                         help="""
+The name of the first scale constant in the `.info` dictionary.
+""".strip()),
+                Argument(name='scale 2', type='float', default=None,
+                         help="""
+A float value for the second scale constant.
+""".strip()),
+                Argument(name='scale 2 name', type='string', default=None,
+                         help="""
+The name of the second scale constant in the `.info` dictionary.
+""".strip()),
+                Argument(name='constant', type='float', default=None,
+                         help="""
+A float value for the offset constant.
+""".strip()),
+                Argument(name='constant name', type='string', default=None,
+                         help="""
+The name of the offset constant in the `.info` dictionary.
+""".strip()),
+                ],
+            help=self.__doc__, plugin=plugin)
+
+    def _run(self, hooke, inqueue, outqueue, params):
+        self._add_to_command_stack(params)
+        i1 = self._get_column(hooke=hooke, params=params,
+                                    column_name='input column 1')
+        i2 = self._get_column(hooke=hooke, params=params,
+                                    column_name='input column 2')
+        if i1 is None:
+            i1 = 0
+        if i2 is None:
+            i2 = 0
+        # what if both i1 and i2 are None?
+        a = self._get_constant(params, i1.info, 'scale 1')
+        b = self._get_constant(params, i2.info, 'scale 2')
+        c = self._get_constant(params, i1.info, 'constant')
+        out = a*i1 + b*i2 + c
+        i1.tofile('1.dat', sep='\n', format='%r')
+        i2.tofile('2.dat', sep='\n', format='%r')
+        out.tofile('3.dat', sep='\n', format='%r')
+        self._set_column(hooke=hooke, params=params,
+                         column_name='output column', values=out)
+
+    def _get_constant(self, params, info, name):
+        a = params[name]
+        pname = params[name + ' name']
+        b = None
+        if pname is not None:
+            pname_entries = pname.split('|')
+            b = info
+            for entry in pname_entries:
+                b = b[entry]
+        if a is None and b is None:
+            return 0
+        if a is None:
+            a = 1
+        if b is None:
+            b = 1
+        return a*b
+
+
 class ClearStackCommand (CurveCommand):
     """Empty a curve's command stack.
     """