Merge pull request #113 from fredrikhellman/force_points_for_one_segment

Add support for force points when a single segment is to be found

pwlf/pwlf.py

@@ -282,6 +282,13 @@ def _fit_one_segment(self):
         """
         self.fit_with_breaks([self.break_0, self.break_n])
 
+    def _fit_one_segment_force_points(self, x_c, y_c):
+        r"""
+        Fit for a single line segment with force points
+        """
+        self.fit_with_breaks_force_points([self.break_0, self.break_n],
+                                          x_c, y_c)
+
     def assemble_regression_matrix(self, breaks, x):
         r"""
         Assemble the linear regression matrix A
@@ -757,7 +764,10 @@ def fit(self, n_segments, x_c=None, y_c=None, bounds=None, **kwargs):
 
         # special fit for one line segment
         if self.n_segments == 1:
-            self._fit_one_segment()
+            if x_c is None and y_c is None:
+                self._fit_one_segment()
+            else:
+                self._fit_one_segment_force_points(self.x_c, self.y_c)
             return self.fit_breaks
 
         # initiate the bounds of the optimization

tests/tests.py

@@ -225,6 +225,28 @@ def test_opt_fit_with_force_points(self):
         yhat = my_fit.predict(x_c)
         self.assertTrue(np.isclose(y_c, yhat))
 
+    def test_opt_fit_single_segment(self):
+        # Test fit for a single segment without force points
+        x = np.linspace(0.0, 1.0, num=100)
+        y = x + 1
+        my_fit = pwlf.PiecewiseLinFit(x, y, disp_res=True)
+        my_fit.fit(1)
+        xhat = 0
+        yhat = my_fit.predict(xhat)
+        self.assertTrue(np.isclose(xhat + 1, yhat))
+
+    def test_opt_fit_with_force_points_single_segment(self):
+        # Test fit for a single segment (same as above)
+        # but with a force point
+        x = np.linspace(0.0, 1.0, num=100)
+        y = x + 1
+        my_fit = pwlf.PiecewiseLinFit(x, y, disp_res=True)
+        x_c = [0.0]
+        y_c = [0.0]
+        my_fit.fit(1, x_c, y_c)
+        yhat = my_fit.predict(x_c)
+        self.assertTrue(np.isclose(y_c, yhat))
+
     def test_se(self):
         # check to see if it will let me calculate standard errors
         my_pwlf = pwlf.PiecewiseLinFit(np.random.random(20),