If AKS selects no internal knots, the model degenerates to a single common intercept across all time points

PiperOrigin-RevId: 865082316

Author: riemanli

meridian/data/test_utils.py

@@ -1953,13 +1953,12 @@ def sample_input_data_for_aks_with_expected_knot_info() -> (
       ),
       'non_revenue',
   )
+  expected_knots = np.array(
+      [0, 11, 14, 38, 39, 41, 43, 45, 48, 50, 55, 87, 89, 90, 116]
+  )
   expected_knot_info = knots.KnotInfo(
-      n_knots=13,
-      knot_locations=np.array(
-          [11, 14, 38, 39, 41, 43, 45, 48, 50, 55, 87, 89, 90]
-      ),
-      weights=knots.l1_distance_weights(
-          117, np.array([11, 14, 38, 39, 41, 43, 45, 48, 50, 55, 87, 89, 90])
-      ),
+      n_knots=15,
+      knot_locations=expected_knots,
+      weights=knots.l1_distance_weights(117, expected_knots),
   )
   return data, expected_knot_info

meridian/model/knots.py

@@ -202,10 +202,9 @@ def get_knot_info(
       )
     else:
       aks = AKS(data)
-      knots = aks.automatic_knot_selection().knots
-      n_knots = len(knots)
-      knot_locations = knots
-  elif isinstance(knots, int):
+      selected_knots = aks.automatic_knot_selection().knots
+      knots = selected_knots if selected_knots.size > 0 else None
+  if isinstance(knots, int):
     if knots < 1:
       raise ValueError('If knots is an integer, it must be at least 1.')
     elif knots > n_times:
@@ -220,7 +219,7 @@ def get_knot_info(
       )
     n_knots = knots
     knot_locations = _get_equally_spaced_knot_locations(n_times, n_knots)
-  elif isinstance(knots, Collection) and knots:
+  elif isinstance(knots, Collection) and len(knots) > 0:
     if any(k < 0 for k in knots):
       raise ValueError('Knots must be all non-negative.')
     if any(k >= n_times for k in knots):
@@ -278,8 +277,11 @@ def automatic_knot_selection(
         value will be used.
 
     Returns:
-      Selected knots and the corresponding B-spline model.
+      Selected knots and the corresponding B-spline model. If at least one knot
+      is selected, boundary knots (min and max time) are added to ensure full
+      time coverage.
     """
+
     if base_penalty is None:
       base_penalty = self._BASE_PENALTY
     n_times = len(self._data.time)
@@ -326,7 +328,15 @@ def automatic_knot_selection(
         np.where(information_criterion == min(information_criterion))[0]
     )
 
-    return AKSResult(knots_sel[opt_idx], model[opt_idx])
+    selected_knots = knots_sel[opt_idx]
+    if selected_knots.size > 0:
+      start_knot = int(x.min())
+      end_knot = int(x.max())
+      selected_knots = np.unique(
+          np.concatenate((selected_knots, [start_knot, end_knot]))
+      )
+
+    return AKSResult(selected_knots, model[opt_idx])
 
   def _get_bspline_matrix(self, x, knots):
     """Replaces patsy.highlevel.dmatrix('bs(...)', ...)"""

meridian/model/knots_test.py

@@ -258,6 +258,80 @@ def test_wrong_knots_fails(
     ):
       knots.get_knot_info(n_times=200, knots=knots_arg, is_national=is_national)
 
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="geo",
+          is_national=False,
+          n_times=5,
+          expected_n_knots=5,
+          expected_knot_locations=[0, 1, 2, 3, 4],
+      ),
+      dict(
+          testcase_name="national",
+          is_national=True,
+          n_times=5,
+          expected_n_knots=1,
+          expected_knot_locations=[0],
+      ),
+  )
+  def test_get_knot_info_aks_returns_empty_falls_back_to_defaults(
+      self,
+      is_national,
+      n_times,
+      expected_n_knots,
+      expected_knot_locations,
+  ):
+    """Tests that if AKS returns empty knots, we fall back to default logic."""
+    mock_result = mock.create_autospec(knots.AKSResult, instance=True)
+    mock_result.knots = np.array([], dtype=int)
+    mock_aks = self.enter_context(
+        mock.patch.object(
+            knots.AKS, "automatic_knot_selection", autospec=True, spec_set=True
+        )
+    )
+    mock_aks.return_value = mock_result
+
+    info = knots.get_knot_info(
+        n_times=n_times,
+        knots=None,
+        enable_aks=True,
+        data=mock.create_autospec(
+            input_data.InputData, instance=True, spec_set=True
+        ),
+        is_national=is_national,
+    )
+
+    self.assertEqual(info.n_knots, expected_n_knots)
+    np.testing.assert_array_equal(info.knot_locations, expected_knot_locations)
+
+  def test_get_knot_info_aks_returns_knots_uses_them(self):
+    mock_result = mock.create_autospec(
+        knots.AKSResult,
+        instance=True,
+    )
+    mock_result.knots = np.array([2, 4], dtype=int)
+    mock_aks = self.enter_context(
+        mock.patch.object(
+            knots.AKS,
+            "automatic_knot_selection",
+            autospec=True,
+            spec_set=True,
+        )
+    )
+    mock_aks.return_value = mock_result
+
+    info = knots.get_knot_info(
+        n_times=10,
+        knots=None,
+        enable_aks=True,
+        data=mock.create_autospec(
+            input_data.InputData, instance=True, spec_set=True
+        ),
+    )
+
+    self.assertEqual(info.n_knots, 2)
+    np.testing.assert_array_equal(info.knot_locations, [2, 4])
+
 
 class AKSTest(parameterized.TestCase):
   """Tests for knots.AKS class."""
@@ -389,7 +463,23 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               ),
               "non_revenue",
           ),
-          expected_knots=[11, 14, 38, 39, 41, 43, 45, 48, 50, 55, 87, 89, 90],
+          expected_knots=[
+              0,
+              11,
+              14,
+              38,
+              39,
+              41,
+              43,
+              45,
+              48,
+              50,
+              55,
+              87,
+              89,
+              90,
+              116,
+          ],
       ),
       dict(
           testcase_name="national_geos",
@@ -404,6 +494,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               "non_revenue",
           ),
           expected_knots=[
+              0,
               1,
               2,
               3,
@@ -457,6 +548,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               103,
               104,
               114,
+              116,
           ],
       ),
       dict(
@@ -472,6 +564,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               "non_revenue",
           ),
           expected_knots=[
+              0,
               4,
               17,
               20,
@@ -490,6 +583,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               77,
               78,
               81,
+              116,
           ],
       ),
       dict(
@@ -504,7 +598,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               ),
               "non_revenue",
           ),
-          expected_knots=[2, 7, 24, 25, 38, 39, 49, 114],
+          expected_knots=[0, 2, 7, 24, 25, 38, 39, 49, 114, 116],
       ),
       dict(
           testcase_name="50_times",
@@ -518,7 +612,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               ),
               "non_revenue",
           ),
-          expected_knots=[1, 5, 13, 15, 16, 23, 27, 31, 32, 38, 42],
+          expected_knots=[0, 1, 5, 13, 15, 16, 23, 27, 31, 32, 38, 42, 49],
       ),
       dict(
           testcase_name="200_times",
@@ -533,6 +627,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               "non_revenue",
           ),
           expected_knots=[
+              0,
               4,
               10,
               12,
@@ -579,6 +674,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               195,
               196,
               197,
+              199,
           ],
       ),
       dict(
@@ -594,7 +690,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               ),
               "non_revenue",
           ),
-          expected_knots=[17, 25],
+          expected_knots=[0, 17, 25, 49],
       ),
       dict(
           testcase_name="seasonal",
@@ -610,6 +706,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               "non_revenue",
           ),
           expected_knots=[
+              0,
               1,
               4,
               5,
@@ -629,6 +726,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               41,
               45,
               47,
+              49,
           ],
       ),
       dict(
@@ -644,7 +742,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               ),
               "non_revenue",
           ),
-          expected_knots=[24, 25, 26],
+          expected_knots=[0, 24, 25, 26, 49],
       ),
       dict(
           testcase_name="minimal_initial_knots",
@@ -658,7 +756,7 @@ def test_aks_internal_knots_guardrail_raises(self, data, expected_error):
               ),
               "non_revenue",
           ),
-          expected_knots=[3],
+          expected_knots=[0, 3, 14],
       ),
   )
   def test_aks(self, data: input_data.InputData, expected_knots: list[int]):
@@ -782,6 +880,7 @@ def test_user_provided_base_penalty(self):
     self.assertListEqual(
         actual_knots.tolist(),
         [
+            0,
             2,
             7,
             8,
@@ -841,6 +940,7 @@ def test_user_provided_base_penalty(self):
             110,
             113,
             114,
+            116,
         ],
     )
 
@@ -849,13 +949,13 @@ def test_user_provided_base_penalty(self):
           testcase_name="min_equals_max",
           min_internal_knots=8,
           max_internal_knots=8,
-          expected_knots=[2, 7, 24, 25, 38, 39, 49, 114],
+          expected_knots=[0, 2, 7, 24, 25, 38, 39, 49, 114, 116],
       ),
       dict(
           testcase_name="min_lt_max_",
           min_internal_knots=2,
           max_internal_knots=15,
-          expected_knots=[2, 7, 24, 25, 38, 39, 49, 114],
+          expected_knots=[0, 2, 7, 24, 25, 38, 39, 49, 114, 116],
       ),
   )
   def test_aks_user_provided_min_max_internal_knots(