Merge pull request #56 from QTC-UMD/parameter_checks

Scannable Parameter checks

Author: dihm

docs/source/changelog.rst

@@ -14,6 +14,8 @@ Improvements
   density matrices into the complex basis with ground state present.
 - Made project `uv` compatible and added installation docs describing how to use it.
 - Updated license metadata to follow PEP 639 standard
+- Improve scannable parameter handling so that all sequences are saved to the graph as numpy arrays.
+  Also ensures proper handling of length-1 and length-0 sequences.
 
 Bug Fixes
 +++++++++

src/rydiqule/sensor.py

@@ -9,7 +9,7 @@
 import itertools
 
 from .sensor_utils import (ScannableParameter, CouplingDict, State, States, StateSpec, StateSpecs, TimeFunc,
-                           match_states, _squeeze_dims, expand_statespec, state_tuple_to_str)
+                           match_states, _squeeze_dims, expand_statespec, state_tuple_to_str, process_scannable_parameter)
 from .exceptions import RydiquleError, CouplingNotAllowedError
 from .exceptions import RWAWarning, PopulationNotConservedWarning, RydiquleWarning, debug_state
 
@@ -416,6 +416,7 @@ def add_single_energy_shift(self, state: State, shift: ScannableParameter, label
             raise RydiquleError(f"state {state} is not a node on the graph")
 
         self._remove_edge_data((state, state), kind="coherent")
+        shift = process_scannable_parameter(shift)
         self.couplings.add_edge(state, state, e_shift=shift, label=label)
         if debug_state():
             print(f'   Added energy shift for {state}')
@@ -856,8 +857,8 @@ def add_single_coupling(
         field_params_trimmed = {k:v for k,v in field_params.items() if v is not None}
 
         full_edge_data = {
-            param: np.array(val)
-            if param in self.scannable_parameters and hasattr(val, "__len__")
+            param: process_scannable_parameter(val)
+            if param in self.scannable_parameters
             else val
             for (param, val) in {**field_params_trimmed, **extra_kwargs}.items()
         }
@@ -1718,8 +1719,7 @@ def add_single_decoherence(self, states: States, gamma: ScannableParameter,
 
         states = self._states_valid(states)
         # coerce gamma to numpy array if a sequence
-        if isinstance(gamma, Sized):
-            gamma = np.array(gamma)
+        gamma = process_scannable_parameter(gamma)
 
         gamma_full = decoherent_cc*gamma
         if np.all(gamma_full==0.0):

src/rydiqule/sensor_utils.py

@@ -11,7 +11,7 @@
 import scipy.constants
 from scipy.constants import hbar, e
 
-from typing import Dict, Tuple, Union, List, Callable, TYPE_CHECKING
+from typing import Dict, Tuple, Union, List, Callable, Sequence, TYPE_CHECKING
 if TYPE_CHECKING:
     # only import when type checking, avoid circular import
     from .sensor import Sensor
@@ -1036,6 +1036,37 @@ def state_tuple_to_str(states:States) -> str:
     return "(" + ",".join([str(s) for s in states]) + ")"
 
 
+def process_scannable_parameter(val: ScannableParameter) -> Union[np.ndarray, float]:
+    """Ensures that scannable parameters are coerced to numpy arrays.
+    
+    If the parameter only has length of one, content is extracted.
+
+    Parameters
+    ----------
+    val : ScannableParameter
+        Scannable parameter to process.
+
+    Returns
+    -------
+    float or numpy.ndarray
+
+    Raises
+    ------
+    RydiquleError:
+        Raised if passed an empty sequence
+    """
+
+    if isinstance(val, (Sequence, np.ndarray)):
+        if len(val) > 1:
+            return np.array(val)
+        elif len(val) == 1:
+            return val[0]
+        else:
+            raise RydiquleError('Length-0 sequence passed as scannable parameter')
+    else:
+        return val
+
+
 def _validate_sols(sols) -> np.ndarray:
     """Helper function to validate that solutions are of an appropriate type.
     There are 3 outcomes: