More tolerant scalar check (#348)

* more tolerant scalar check

* refactor broadcast of value/lower/upper etc.

* isort

* change allow_none to default False

* add missing backtick, good job flake8

* version bump

Author: ewu63

pyoptsparse/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.10.0"
+__version__ = "2.10.1"
 
 from .pyOpt_history import History
 from .pyOpt_variable import Variable

pyoptsparse/pyOpt_constraint.py

@@ -8,7 +8,7 @@
 
 # Local modules
 from .pyOpt_error import Error, pyOptSparseWarning
-from .pyOpt_utils import INFINITY, convertToCOO
+from .pyOpt_utils import INFINITY, _broadcast_to_array, convertToCOO
 from .types import Dict1DType
 
 
@@ -43,41 +43,9 @@ def __init__(
         # Before we can do the processing below we need to have lower
         # and upper arguments expanded:
 
-        if lower is None:
-            lower = [None for i in range(self.ncon)]
-        elif np.isscalar(lower):
-            lower = lower * np.ones(self.ncon)
-        elif len(lower) == self.ncon:
-            pass  # Some iterable object
-        else:
-            raise Error(
-                "The 'lower' argument to addCon or addConGroup is invalid. "
-                + f"It must be None, a scalar, or a list/array or length nCon={nCon}."
-            )
-
-        if upper is None:
-            upper = [None for i in range(self.ncon)]
-        elif np.isscalar(upper):
-            upper = upper * np.ones(self.ncon)
-        elif len(upper) == self.ncon:
-            pass  # Some iterable object
-        else:
-            raise Error(
-                "The 'upper' argument to addCon or addConGroup is invalid. "
-                + f"It must be None, a scalar, or a list/array or length nCon={nCon}."
-            )
-
-        # ------ Process the scale argument
-        scale = np.atleast_1d(scale)
-        if len(scale) == 1:
-            scale = scale[0] * np.ones(nCon)
-        elif len(scale) == nCon:
-            pass
-        else:
-            raise Error(
-                f"The length of the 'scale' argument to addCon or addConGroup is {len(scale)}, "
-                + f"but the number of constraints is {nCon}."
-            )
+        lower = _broadcast_to_array("lower", lower, nCon, allow_none=True)
+        upper = _broadcast_to_array("upper", upper, nCon, allow_none=True)
+        scale = _broadcast_to_array("scale", scale, nCon)
 
         # Save lower and upper...they are only used for printing however
         self.lower = lower

pyoptsparse/pyOpt_optimization.py

@@ -15,7 +15,18 @@
 from .pyOpt_constraint import Constraint
 from .pyOpt_error import Error
 from .pyOpt_objective import Objective
-from .pyOpt_utils import ICOL, IDATA, INFINITY, IROW, convertToCOO, convertToCSR, mapToCSR, scaleColumns, scaleRows
+from .pyOpt_utils import (
+    ICOL,
+    IDATA,
+    INFINITY,
+    IROW,
+    _broadcast_to_array,
+    convertToCOO,
+    convertToCSR,
+    mapToCSR,
+    scaleColumns,
+    scaleRows,
+)
 from .pyOpt_variable import Variable
 from .types import Dict1DType, Dict2DType, NumpyType
 
@@ -229,71 +240,11 @@ def addVarGroup(
         if varType not in ["c", "i", "d"]:
             raise Error("Type must be one of 'c' for continuous, 'i' for integer or 'd' for discrete.")
 
-        # ------ Process the value argument
-        value = np.atleast_1d(value).real
-        if len(value) == 1:
-            value = value[0] * np.ones(nVars)
-        elif len(value) == nVars:
-            pass
-        else:
-            raise Error(
-                f"The length of the 'value' argument to addVarGroup is {len(value)}, "
-                + f"but the number of variables in nVars is {nVars}."
-            )
-
-        if lower is None:
-            lower = [None for i in range(nVars)]
-        elif np.isscalar(lower):
-            lower = lower * np.ones(nVars)
-        elif len(lower) == nVars:
-            lower = np.atleast_1d(lower).real
-        else:
-            raise Error(
-                "The 'lower' argument to addVarGroup is invalid. "
-                + f"It must be None, a scalar, or a list/array or length nVars={nVars}."
-            )
-
-        if upper is None:
-            upper = [None for i in range(nVars)]
-        elif np.isscalar(upper):
-            upper = upper * np.ones(nVars)
-        elif len(upper) == nVars:
-            upper = np.atleast_1d(upper).real
-        else:
-            raise Error(
-                "The 'upper' argument to addVarGroup is invalid. "
-                + f"It must be None, a scalar, or a list/array or length nVars={nVars}."
-            )
-
-        # ------ Process the scale argument
-        if scale is None:
-            scale = np.ones(nVars)
-        else:
-            scale = np.atleast_1d(scale)
-            if len(scale) == 1:
-                scale = scale[0] * np.ones(nVars)
-            elif len(scale) == nVars:
-                pass
-            else:
-                raise Error(
-                    f"The length of the 'scale' argument to addVarGroup is {len(scale)}, "
-                    + f"but the number of variables in nVars is {nVars}."
-                )
-
-        # ------ Process the offset argument
-        if offset is None:
-            offset = np.ones(nVars)
-        else:
-            offset = np.atleast_1d(offset)
-            if len(offset) == 1:
-                offset = offset[0] * np.ones(nVars)
-            elif len(offset) == nVars:
-                pass
-            else:
-                raise Error(
-                    f"The length of the 'offset' argument to addVarGroup is {len(offset)}, "
-                    + f"but the number of variables is {nVars}."
-                )
+        value = _broadcast_to_array("value", value, nVars)
+        lower = _broadcast_to_array("lower", lower, nVars, allow_none=True)
+        upper = _broadcast_to_array("upper", upper, nVars, allow_none=True)
+        scale = _broadcast_to_array("scale", scale, nVars)
+        offset = _broadcast_to_array("offset", offset, nVars)
 
         # Determine if scalar i.e. it was called from addVar():
         scalar = kwargs.pop("scalar", False)

pyoptsparse/pyOpt_utils.py

@@ -20,6 +20,7 @@
 
 # Local modules
 from .pyOpt_error import Error
+from .types import ArrayType
 
 # Define index mnemonics
 IROW = 0
@@ -529,3 +530,43 @@ def _csc_to_coo(mat: dict) -> dict:
     coo_data = np.array(data)
 
     return {"coo": [coo_rows, coo_cols, coo_data], "shape": mat["shape"]}
+
+
+def _broadcast_to_array(name: str, value: ArrayType, n_values: int, allow_none: bool = False):
+    """
+    Broadcast an input to an array with a specified length
+
+    Parameters
+    ----------
+    name : str
+        The name of the input. This is only used in the error message emitted.
+    value : float, list[float], numpy array
+        The input value
+    n_values : int
+        The number of values
+    allow_none : bool, optional
+        Whether to allow `None` in the input/output, by default False
+
+    Returns
+    -------
+    NDArray
+        An array with the shape ``(n_values)``
+
+    Raises
+    ------
+    Error
+        If either the input is not broadcastable, or if the input contains None and ``allow_none=False``.
+
+    Warnings
+    --------
+    Note that the default value for ``allow_none`` is False.
+    """
+    try:
+        value = np.broadcast_to(value, n_values)
+    except ValueError:
+        raise Error(
+            f"The '{name}' argument is invalid. It must be None, a scalar, or a list/array or length {n_values}."
+        )
+    if not allow_none and any([i is None for i in value]):
+        raise Error(f"The {name} argument cannot be 'None'.")
+    return value

pyoptsparse/types.py

@@ -1,14 +1,15 @@
 # Standard Python modules
-from typing import Dict, List, Union
+from typing import Dict, Sequence, Union
 
 # External modules
-from numpy import ndarray
+import numpy as np
+import numpy.typing as npt
 
 # Either ndarray or scalar
-NumpyType = Union[float, ndarray]
+NumpyType = Union[float, npt.NDArray[np.float_]]
 # ndarray, list of numbers, or scalar
-ArrayType = Union[float, List[float], ndarray]
+ArrayType = Union[NumpyType, Sequence[float]]
 # funcs
-Dict1DType = Dict[str, ndarray]
+Dict1DType = Dict[str, npt.NDArray[np.float_]]
 # funcsSens
-Dict2DType = Dict[str, Dict[str, ndarray]]
+Dict2DType = Dict[str, Dict[str, npt.NDArray[np.float_]]]