Allow make_broadcastable to preserve old function call behaviour

Author: willGraham01

movement/utils/broadcasting.py

@@ -174,7 +174,9 @@ def make_broadcastable(  # noqa: C901
     where the ``self`` argument is present only if ``f`` was a class method.
     ``fr`` applies ``f`` along the ``broadcast_dimension`` of ``data``.
     The ``\*args`` and ``\*\*kwargs`` match those passed to ``f``, and retain
-    the same interpretations and effects on the result.
+    the same interpretations and effects on the result. If ``data`` provided to
+    ``fr`` is not an ``xarray.DataArray``, it will fall back on the behaviour
+    of ``f`` (and ignore the ``broadcast_dimension`` argument).
 
     See the docstring of ``make_broadcastable_inner`` in the source code for a
     more explicit explanation of the returned decorator.
@@ -268,6 +270,9 @@ def inner_clsmethod(  # type: ignore[valid-type]
             broadcast_dimension: str = "space",
             **kwargs: KeywordArgs.kwargs,
         ) -> xr.DataArray:
+            # Preserve original functionality
+            if not isinstance(data, xr.DataArray):
+                return f(self, data, *args, **kwargs)
             return apply_along_da_axis(
                 lambda input_1D: f(self, input_1D, *args, **kwargs),
                 data,
@@ -297,6 +302,9 @@ def inner(  # type: ignore[valid-type]
             broadcast_dimension: str = "space",
             **kwargs: KeywordArgs.kwargs,
         ) -> xr.DataArray:
+            # Preserve original functionality
+            if not isinstance(data, xr.DataArray):
+                return f(data, *args, **kwargs)
             return apply_along_da_axis(
                 lambda input_1D: f(input_1D, *args, **kwargs),
                 data,

tests/test_unit/test_make_broadcastable.py

@@ -241,3 +241,25 @@ def two_to_some(xy_pair) -> np.ndarray:
         else:
             assert d in mock_dataset.dims
             assert len(output[d]) == len(mock_dataset[d])
+
+
+def test_retain_underlying_function() -> None:
+    value_for_arg = 5.0
+    value_for_kwarg = 7.0
+    value_for_simple_input = [0.0, 1.0, 2.0]
+
+    def simple_function(input_1D, arg, kwarg=3.0):
+        return arg * sum(input_1D) + kwarg
+
+    @make_broadcastable()
+    def simple_function_broadcastable(input_1D, arg, kwarg=3.0):
+        return simple_function(input_1D, arg, kwarg=kwarg)
+
+    result_from_broadcastable = simple_function_broadcastable(
+        value_for_simple_input, value_for_arg, kwarg=value_for_kwarg
+    )
+    result_from_original = simple_function(
+        value_for_simple_input, value_for_arg, kwarg=value_for_kwarg
+    )
+    assert isinstance(result_from_broadcastable, float)
+    assert np.isclose(result_from_broadcastable, result_from_original)