Implementing scanning of coupled parameters such as Surface RoC

playground/homs/roc_scan_nomatrix.py

@@ -23,10 +23,18 @@ cp cav_w FP w
 #bp cav_q q ITM.p2.o
 #bp cav_w w ITM.p2.o
 
-#xaxis ETM.Rc lin 2.5 10 200
-xaxis CAV.L lin 1 6 100
+#xaxis ETM.Rc lin 2.5 0.95 200
+#xaxis CAV.L lin 1 6 100
 """
 )
 
-out = model.run()
+# out = finesse.analysis.xaxis(model.ETM.Rc, "lin", 2.5, 0.95, 200)
+out = finesse.analysis.x2axis(
+    model.ETM.Rc, "lin", 2.5, 10, 100, model.ITM.Rc, "lin", -2.5, -10, 100,
+)
+
+print(model.ETM.Rc)
+print(model.ITM.Rc)
+
+# out = model.run()
 out.plot()

src/finesse/analysis/actions.py

@@ -5,6 +5,7 @@ from finesse.solutions import BaseSolution
 from finesse.analysis.runners import run_axes_scan
 from finesse.solutions import ArraySolution
 from finesse.element import ParameterLocked, Parameter
+from finesse.utilities.misc import is_iterable
 
 import regex as re
 import logging
@@ -177,7 +178,15 @@ class StepParamNDWorkspace(ActionWorkspace):
 
 
 class StepParamND(Action):
-    def __init__(self, name, *args, pre_step=None, post_step=None, on_complete=None):
+    def __init__(
+        self,
+        name,
+        *args,
+        pre_step=None,
+        post_step=None,
+        on_complete=None,
+        reset_values=False,
+    ):
         super().__init__(name)
 
         if len(args) % 3 != 0:
@@ -186,14 +195,65 @@ class StepParamND(Action):
                 "values to scan over, and offset to array values."
             )
 
-        self.args = args
+        new_args = []
+        self.initial_parameter_values = {}
+        self.__reset = reset_values
+        if args:
+            first_param = args[0]
+            if isinstance(first_param, Parameter):
+                model = first_param.owner._model
+            else:
+                if not is_iterable(first_param) or not isinstance(
+                    first_param[0], Parameter
+                ):
+                    raise TypeError(
+                        f"Unexpected argument type for parameter: {first_param}"
+                    )
+
+                model = first_param[0].owner._model
+
+            for i, p in enumerate(args[::3]):
+                if isinstance(p, Parameter):
+                    if p.owner._model != model:
+                        raise ValueError(
+                            f"Parameter {p} is not from the same model as "
+                            "the other specified parameters."
+                        )
+
+                    self.initial_parameter_values[p] = p.value
+                    new_args.extend(args[i * 3 : (i * 3) + 3])
+                else:
+                    if not is_iterable(p):
+                        raise TypeError(f"Unexpected argument type for parameter: {p}")
+
+                    if any(p_.owner._model != model for p_ in p):
+                        raise ValueError(
+                            f"Coupled parameters {p} are not from the same model as "
+                            "the other specified parameters."
+                        )
+
+                    # Set up temporary references between the coupled parameters
+                    # -> all parameters get set to ref of first param in p
+                    source_param_ref = p[0].ref
+                    self.initial_parameter_values[p[0]] = p[0].value
+                    new_args.append(p[0])
+                    new_args.extend(args[i * 3 + 1 : (i * 3) + 3])
+                    for coupled_param in p[1:]:
+                        self.initial_parameter_values[
+                            coupled_param
+                        ] = coupled_param.value
+                        setattr(
+                            coupled_param.owner, coupled_param.name, source_param_ref
+                        )
+
+        self.args = new_args
         self._info.parameters_changing = tuple(
             f"{p.component.name}.{p.name}" if type(p) is Parameter else p
-            for p in args[::3]
+            for p in self.args[::3]
         )
 
-        self.axes = tuple(np.atleast_1d(_) for _ in args[1::3])
-        self.offsets = np.array(args[2::3], dtype=np.float64)
+        self.axes = tuple(np.atleast_1d(_) for _ in self.args[1::3])
+        self.offsets = np.array(self.args[2::3], dtype=np.float64)
         self.out_shape = tuple(np.size(_) for _ in self.axes)
         self._info.makes_solution = True
 
@@ -268,6 +328,12 @@ class StepParamND(Action):
             ws = Folder("on_complete", self.on_complete).setup(ws.sol, ws.model)
             ws.fn_do(ws)
 
+        if self.__reset:
+            # Reset each model parameter back to its starting value
+            for p, value in self.initial_parameter_values.items():
+                p._locked = False  # bit of a hack but don't think it matters
+                setattr(p.owner, p.name, value)
+
 
 class Noxaxis(StepParamND):
     def __init__(self, *, name="noxaxis", **kwargs):
@@ -304,6 +370,7 @@ class XNaxis(StepParamND):
             pre_step=pre_step,
             post_step=post_step,
             on_complete=on_complete,
+            reset_values=True,  # xN-axis analyses should always reset model parameters on completion
         )
 
     def setup(self, s_prev, model):

src/finesse/analysis/axes.py

@@ -3,11 +3,26 @@
 import logging
 
 from finesse.element import Parameter
+from finesse.utilities.misc import is_iterable
 import finesse.analysis.actions as ac
 
 LOGGER = logging.getLogger(__name__)
 
 
+def _get_model(param):
+    if is_iterable(param):
+        if any(not isinstance(p, Parameter) for p in param):
+            raise ValueError(
+                f"Expected all scanned objects in {param} to " "be of type Parameter."
+            )
+
+        p = param[0]
+    else:
+        p = param
+
+    return p.owner._model
+
+
 def noxaxis(model):
     analysis = ac.StepParamND("noxaxis")
     return analysis.run(model).get(analysis.name)
@@ -40,10 +55,9 @@ def xaxis(
     offset : float, optional
         Offset to scanned values. For a given xaxis point, `param` will be set to `x[i] + offset`.
     """
-    analysis = ac.Xaxis(
-        f"{param.owner.name}.{param.name}", mode, start, stop, steps, offset=offset
-    )
-    return analysis.run(param.owner._model).get(analysis.name)
+    model = _get_model(param)
+    analysis = ac.Xaxis(param, mode, start, stop, steps, offset=offset)
+    return analysis.run(model).get(analysis.name)
 
 
 def x2axis(
@@ -83,18 +97,20 @@ def x2axis(
     -----
     `param2` is in the inner loop and `param1` in the outer loop.
     """
-    if param1.owner._model is not param2.owner._model:
+    model1 = _get_model(param1)
+    model2 = _get_model(param2)
+    if model1 is not model2:
         raise ValueError(
             f"Parameters {param1} and {param2}, are from different models."
         )
 
     analysis = ac.X2axis(
-        f"{param1.owner.name}.{param1.name}",
+        param1,
         mode1,
         start1,
         stop1,
         steps1,
-        f"{param2.owner.name}.{param2.name}",
+        param2,
         mode2,
         start2,
         stop2,
@@ -103,7 +119,7 @@ def x2axis(
         offset2=offset2,
     )
 
-    return analysis.run(param1.owner._model).get(analysis.name)
+    return analysis.run(model1).get(analysis.name)
 
 
 def x3axis(
@@ -149,22 +165,26 @@ def x3axis(
     -----
     `param3` is in the inner loop and `param1` in the outer loop.
     """
-    if (param1.owner._model is not param2.owner._model) is not param3.owner._model:
+    model1 = _get_model(param1)
+    model2 = _get_model(param2)
+    model3 = _get_model(param3)
+    if model1 is not model2 is not model3:
         raise ValueError(
-            f"Parameters {param1}, {param2}, and {param3}, are not from the same models."
+            f"Parameters {param1}, {param2}, and {param3}, are from different models."
         )
+
     analysis = ac.X3axis(
-        f"{param1.owner.name}.{param1.name}",
+        param1,
         mode1,
         start1,
         stop1,
         steps1,
-        f"{param2.owner.name}.{param2.name}",
+        param2,
         mode2,
         start2,
         stop2,
         steps2,
-        f"{param3.owner.name}.{param3.name}",
+        param3,
         mode3,
         start3,
         stop3,
@@ -174,4 +194,4 @@ def x3axis(
         offset3=offset3,
     )
 
-    return analysis.run(param1.owner._model).get(analysis.name)
+    return analysis.run(model1).get(analysis.name)

src/finesse/components/surface.py

@@ -4,7 +4,7 @@ import numpy as np
 
 from finesse.components.general import Connector
 from finesse.element import model_parameter, Rebuild
-from finesse.utilities.misc import calltracker
+from finesse.utilities.misc import calltracker, is_iterable
 
 
 @model_parameter("R", 0.5, Rebuild.PlaneWave, validate="_check_R", setter="set_RTL")
@@ -219,12 +219,10 @@ following must be specified:
 
     @property
     def Rc(self):
-        """The radius of curvature of the mirror in metres. If the
-        mirror is astigmatic this is a tuple of the tangential and
-        sagittal radii of curvature.
+        """The radius of curvature of the mirror for both planes.
 
-        :getter: Returns the radius of curvature of the mirror in metres (a
-                 tuple of values if the mirror is astigmatic).
+        :getter: Returns the radius of curvature of the mirror for both
+                 the tangential and sagittal planes.
 
         :setter: Sets the radius of curvature.
 
@@ -241,21 +239,26 @@ following must be specified:
 
         >>> obj.Rc = (2.5, 2.7)
         """
-        if np.isclose(self.Rcx.value, self.Rcy.value):
-            return self.Rcx
         return (self.Rcx, self.Rcy)
 
     @Rc.setter
     def Rc(self, value):
-        try:
-            self.Rcx = value[0]
-            self.Rcy = value[1]
-        except (IndexError, TypeError):
+        if is_iterable(value):
+            if len(value) != 2:
+                raise ValueError(
+                    f"Error in setting {self.name}.Rc -> expected a single value"
+                    f"or an iterable of two values for x, y planes but got: {value}."
+                )
+
+            Rcx, Rcy = value
+            self._check_Rc(Rcx)
+            self._check_Rc(Rcy)
+            self.Rcx, self.Rcy = Rcx, Rcy
+        else:
+            self._check_Rc(value)
             self.Rcx = value
             self.Rcy = value
 
-        self._check_Rc(value)
-
     # NOTE this is a bit hacky but gets around using surface.R = value (etc.)
     #      directly in an axis scan without being warned
 

src/finesse/utilities/misc.py

@@ -99,6 +99,22 @@ def find_nearest(x, value, index=False):
     return x[idx]
 
 
+def is_iterable(obj):
+    """Reliable check for whether an object is iterable.
+
+    Returns
+    -------
+    flag : bool
+        True if `obj` is iterable, False otherwise.
+    """
+    try:
+        iter(obj)
+    except Exception:
+        return False
+    else:
+        return True
+
+
 def ngettext(n, fsingle, fplural, sub=True):
     """Get the singular or plural form of the specified messages based on n.
 

tests/functional/components/__init__.py

@@ -125,7 +125,8 @@ class TestSurface(TestComponent, metaclass=abc.ABCMeta):
         """Test that setting a surface's Rc to a single value sets Rcx and Rcy to that value."""
         obj = component(name="cmp1")
         obj.Rc = 3.141
-        assert float(obj.Rc) == 3.141
+        Rcx, Rcy = obj.Rc
+        assert (float(Rcx), float(Rcy)) == (3.141, 3.141)
         assert float(obj.Rcx) == 3.141
         assert float(obj.Rcy) == 3.141
 
@@ -139,6 +140,7 @@ class TestSurface(TestComponent, metaclass=abc.ABCMeta):
         assert float(obj.Rcx) == 3.141
         assert float(obj.Rcy) == 6.282
 
+    @pytest.mark.skip(reason="Surface.Rc always returns a tuple on this branch")
     def test_same_rcx_and_rcy_makes_rc_return_single_number(self, component):
         """Test that setting a surface's Rcx and Rcy to the same number makes Rc return a single \
         number."""