From c92fd9b85ac56df6b998fae3f0ff5c93e52c9708 Mon Sep 17 00:00:00 2001
From: Ardavan Oskooi <ardavano@google.com>
Date: Tue, 26 Dec 2023 20:23:06 -0800
Subject: [PATCH] add type checking for frequencies argument of
 OptimizationProblem constructor (#2747)

---
 python/adjoint/optimization_problem.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/python/adjoint/optimization_problem.py b/python/adjoint/optimization_problem.py
index d72d89db0..628c48877 100644
--- a/python/adjoint/optimization_problem.py
+++ b/python/adjoint/optimization_problem.py
@@ -28,7 +28,7 @@ def __init__(
         objective_functions: List[Callable],
         objective_arguments: List[ObjectiveQuantity],
         design_regions: List[DesignRegion],
-        frequencies: Optional[Union[float, List[float]]] = None,
+        frequencies: Optional[Union[np.ndarray, List[float]]] = None,
         fcen: Optional[float] = None,
         df: Optional[float] = None,
         nf: Optional[int] = None,
@@ -53,7 +53,7 @@ def __init__(
             and we have to specify arguments as [A,B,C]
           objective_arguments: list of ObjectiveQuantity passed as arguments of objective functions
           design_regions: list of DesignRegion to be optimized
-          frequencies: list of frequencies of interest in the problem. If not specified, then the list of frequencies
+          frequencies: array/list of frequencies of interest in the problem. If not specified, then the list of frequencies
             will be created from fcen, df, and nf: a list of size nf that goes from fcen-df/2 to fcen+df/2
           fcen: center frequency
           df: range of frequencies, i.e. maximum frequency -  minimum frequency
@@ -88,10 +88,18 @@ def __init__(
         self.num_design_params = [ni.num_design_params for ni in self.design_regions]
         self.num_design_regions = len(self.design_regions)
 
-        # TODO typecheck frequency choices
         if frequencies is not None:
-            self.frequencies = frequencies
-            self.nf = np.array(frequencies).size
+            if isinstance(frequencies, (np.ndarray, List)):
+                self.frequencies = frequencies
+                if isinstance(frequencies, np.ndarray):
+                    self.nf = frequencies.size
+                else:
+                    self.nf = len(frequencies)
+            else:
+                raise TypeError(
+                    "frequencies argument of OptimizationProblem "
+                    "constructor must be a Numpy array or List."
+                )
         elif nf == 1:
             self.nf = nf
             self.frequencies = [fcen]