Be clear that quadratic can be used with a float or an array

src/openflexure_microscope_server/utilities.py

@@ -1,6 +1,15 @@
 """Utility functions and classes."""
 
-from typing import TypeVar, Callable, ParamSpec, Optional, Any, Concatenate, Self
+from typing import (
+    TypeVar,
+    Callable,
+    ParamSpec,
+    Optional,
+    Any,
+    Concatenate,
+    Self,
+    overload,
+)
 import os
 import re
 import sys
@@ -11,6 +20,7 @@ import tomllib
 from functools import wraps
 
 from pydantic import BaseModel
+import numpy as np
 
 T = TypeVar("T")
 P = ParamSpec("P")
@@ -329,10 +339,20 @@ def _get_version_from_toml(toml_path: str) -> str:
         return "Undefined"
 
 
-def quadratic(x: float, a: float, b: float, c: float) -> float:
+# Let MyPy know that the output type matches x.
+@overload
+def quadratic(x: np.ndarray, a: float, b: float, c: float) -> np.ndarray: ...
+@overload
+def quadratic(x: float, a: float, b: float, c: float) -> float: ...
+
+
+def quadratic(
+    x: float | np.ndarray, a: float, b: float, c: float
+) -> float | np.ndarray:
     """Quadratic function. Used for predicting z.
 
-    :param x: The points at which to evaluate the quadratic.
+    :param x: The point or points at which to evaluate the quadratic. This can be a
+        float or a numpy array. The return will be the same type.
     :param a: The coefficient of x^2.
     :param b: The coefficient of x.
     :param c: The constant coefficient.