def equal_length_array_or_scalar(
array, length=1, mode="continue"
):
"""
Returns 'array' if its length is equal to 'length'. If this is not the
case, returns an array of length 'length' with values equal to the first
value of the array (or if 'array' is a scalar, that value. If mode is
"warn", issues a warning if this happens; if mode is "exception" raises an
exception in this case.
"""
try:
array_length = len(array)
if array_length == length:
return array
else:
if mode == "warn":
warnings.warn("Length of array is not equal to %i. Using only\
the first value." % length)
try:
unit = array.unit
value = array[0].value_in(unit)
except:
unit = units.none
value = array[0]
array = VectorQuantity(
array=numpy.ones(length) * value,
unit=unit,
)
return array
elif mode == "exception":
raise Exception("Length of array is not equal to %i. This is\
not supported." % length)
except:
try:
unit = array.unit
value = array.value_in(unit)
except:
unit = units.none
value = array
array = VectorQuantity(
array=numpy.ones(length) * value,
unit=unit,
)
if mode == "warn":
warnings.warn("Using single value for all cases.")
return array
def equal_length_array_or_scalar(array, length=1, mode="continue"):
"""
Returns 'array' if its length is equal to 'length'. If this is not the
case, returns an array of length 'length' with values equal to the first
value of the array (or if 'array' is a scalar, that value. If mode is
"warn", issues a warning if this happens; if mode is "exception" raises an
exception in this case.
"""
try:
array_length = len(array)
if array_length == length:
return array
else:
if mode == "warn":
warnings.warn("Length of array is not equal to %i. Using only\
the first value." % length)
try:
unit = array.unit
value = array[0].value_in(unit)
except:
unit = units.none
value = array[0]
array = VectorQuantity(
array=numpy.ones(length) * value,
unit=unit,
)
return array
elif mode == "exception":
raise Exception("Length of array is not equal to %i. This is\
not supported." % length)
except:
try:
unit = array.unit
value = array.value_in(unit)
except:
unit = units.none
value = array
array = VectorQuantity(
array=numpy.ones(length) * value,
unit=unit,
)
if mode == "warn":
warnings.warn("Using single value for all cases.")
return array
