Repository: wright-group/attune

Issue:
method to add dependent(s) to curve

Discussion:

Patch:
diff --git a/attune/curve/_base.py b/attune/curve/_base.py
index 71d2ae9..e3982a4 100644
--- a/attune/curve/_base.py
+++ b/attune/curve/_base.py
@@ -57,7 +57,12 @@ class Curve:
 
         self.__version__ = __version__
         self.setpoints = setpoints
-        self.dependents = {d.name: d for d in dependents}
+        if isinstance(dependents, Dependent):
+            dependents = [dependents]
+        if isinstance(dependents, dict):
+            self.dependents = dependents
+        else:
+            self.dependents = {d.name: d for d in dependents}
         self.name = name
         self.kind = kind
         self.subcurve = subcurve
@@ -65,7 +70,7 @@ class Curve:
         self.interaction = interaction
         # set dependents as attributes of self
         for obj in self.dependents.values():
-            if len(obj[:]) != len(self.setpoints[:]):
+            if len(obj) != len(self.setpoints):
                 raise ValueError("Dependents must be the same length as setpoints")
             setattr(self, obj.name, obj)
         # initialize function object
@@ -119,6 +124,16 @@ class Curve:
     def __getitem__(self, key):
         return self.dependents[key]
 
+    def __setitem__(self, key, value):
+        value = copy_.deepcopy(value)
+        value.name = key
+        if value.interpolator is not None:
+            value.positions = value(self.setpoints[:], self.setpoints.units)
+        elif len(value) != len(self.setpoints):
+            raise ValueError(f"Incorrect number of points in dependent: {len(value)} for number of setpoints: {len(self.setpoints)}")
+        value.interpolator = self.method(self.setpoints, value)
+        self.dependents[key] = value
+
     def __call__(self, value, units=None, full=True):
         return self.get_dependent_positions(value, units, full=full)
 
@@ -498,7 +513,7 @@ class Curve:
         else:
             save_directory = pathlib.Path(save_directory)
         # array
-        out_arr = np.zeros([len(self.dependents) + 1, len(self.setpoints[:])])
+        out_arr = np.zeros([len(self.dependents) + 1, len(self.setpoints)])
         out_arr[0] = self.setpoints[:]
         out_arr[1:] = np.array([dependent.positions for dependent in self.dependents.values()])
         # filename
diff --git a/attune/curve/_dependent.py b/attune/curve/_dependent.py
index 2e31b92..6c79722 100644
--- a/attune/curve/_dependent.py
+++ b/attune/curve/_dependent.py
@@ -15,6 +15,9 @@ class Variable(object):
     def __setitem__(self, key, value):
         self.positions[key] = value
 
+    def __len__(self):
+        return len(self.positions)
+
     def convert(self, units):
         self.positions = wt.units.convert(self.positions, self.units, units)
         self.units = units


Test patch:
diff --git a/tests/Curve/setitem.py b/tests/Curve/setitem.py
new file mode 100644
index 0000000..e56caa0
--- /dev/null
+++ b/tests/Curve/setitem.py
@@ -0,0 +1,40 @@
+import attune
+import numpy as np
+import pytest
+
+def test_set_new():
+    d = attune.Dependent(np.linspace(0,10,21), "d1")
+    s = attune.Setpoints(np.linspace(120, 240, 21), "s1")
+    c = attune.Curve(s, d, "c")
+    d = attune.Dependent(np.linspace(10,100,21), "d2")
+
+    c["d3"] = d
+
+    assert "d3" in c.dependents
+    assert "d2" not in c.dependents
+    assert np.allclose(c["d3"][:], d[:])
+    assert c["d3"].name == "d3"
+
+def test_set_incompatable():
+    d = attune.Dependent(np.linspace(0,10,21), "d1")
+    s = attune.Setpoints(np.linspace(120, 240, 21), "s1")
+    c = attune.Curve(s, d, "c")
+    d = attune.Dependent(np.linspace(10,100,25), "d2")
+
+    with pytest.raises(ValueError):
+        c["d3"] = d
+
+def test_set_from_outside_curve():
+    d = attune.Dependent(np.linspace(0,10,21), "d1")
+    s = attune.Setpoints(np.linspace(120, 240, 21), "s1")
+    c = attune.Curve(s, d, "c")
+    d = attune.Dependent(np.linspace(10,100,30), "d2")
+    s = attune.Setpoints(np.linspace(100, 250, 30), "s2")
+    c2 = attune.Curve(s, d, "c")
+
+    c["d3"] = c2["d2"]
+
+    assert "d3" in c.dependents
+    assert "d2" not in c.dependents
+    assert np.allclose(c["d3"][:], d(c.setpoints[:]))
+    assert c["d3"].name == "d3"