test/func: modifying refinement to allow lists and dicts and adding a test of refinement for squeeze and funcy

Author: Luiskitsu

src/diffpy/morph/morph_api.py

@@ -39,9 +39,17 @@
         morph_helpers.TransformXtalRDFtoPDF,
     ],
     qdamp=morphs.MorphResolutionDamping,
+    squeeze=morphs.MorphSqueeze,
+    parameters=morphs.MorphFuncy,
 )
 _default_config = dict(
-    scale=None, stretch=None, smear=None, baselineslope=None, qdamp=None
+    scale=None,
+    stretch=None,
+    smear=None,
+    baselineslope=None,
+    qdamp=None,
+    squeeze=None,
+    parameters=None,
 )
 
 
@@ -197,6 +205,14 @@ def morph(
         if k == "smear":
             [chain.append(el()) for el in morph_cls]
             refpars.append("baselineslope")
+        elif k == "parameters":
+            morph_inst = morph_cls()
+            morph_inst.function = rv_cfg.get("function", None)
+            if morph_inst.function is None:
+                raise ValueError(
+                    "Must provide a 'function' when using 'parameters'"
+                )
+            chain.append(morph_inst)
         else:
             chain.append(morph_cls())
         refpars.append(k)

src/diffpy/morph/morphs/__init__.py

@@ -19,13 +19,15 @@
 
 from diffpy.morph.morphs.morph import Morph  # noqa: F401
 from diffpy.morph.morphs.morphchain import MorphChain  # noqa: F401
+from diffpy.morph.morphs.morphfuncy import MorphFuncy
 from diffpy.morph.morphs.morphishape import MorphISphere, MorphISpheroid
 from diffpy.morph.morphs.morphresolution import MorphResolutionDamping
 from diffpy.morph.morphs.morphrgrid import MorphRGrid
 from diffpy.morph.morphs.morphscale import MorphScale
 from diffpy.morph.morphs.morphshape import MorphSphere, MorphSpheroid
 from diffpy.morph.morphs.morphshift import MorphShift
 from diffpy.morph.morphs.morphsmear import MorphSmear
+from diffpy.morph.morphs.morphsqueeze import MorphSqueeze
 from diffpy.morph.morphs.morphstretch import MorphStretch
 
 # List of morphs
@@ -40,6 +42,8 @@
     MorphISpheroid,
     MorphResolutionDamping,
     MorphShift,
+    MorphSqueeze,
+    MorphFuncy,
 ]
 
 # End of file

src/diffpy/morph/morphs/morphfuncy.py

@@ -11,6 +11,7 @@ class MorphFuncy(Morph):
     yinlabel = LABEL_GR
     xoutlabel = LABEL_RA
     youtlabel = LABEL_GR
+    parnames = ["parameters"]
 
     def morph(self, x_morph, y_morph, x_target, y_target):
         """General morph function that applies a user-supplied function to the

src/diffpy/morph/refine.py

@@ -26,8 +26,6 @@
 class Refiner(object):
     """Class for refining a Morph or MorphChain.
 
-    This is provided to allow for custom residuals and refinement algorithms.
-
     Attributes
     ----------
     chain
@@ -51,12 +49,27 @@ def __init__(self, chain, x_morph, y_morph, x_target, y_target):
         self.y_target = y_target
         self.pars = []
         self.residual = self._residual
+        self.flat_to_grouped = {}
         return
 
     def _update_chain(self, pvals):
         """Update the parameters in the chain."""
-        pairs = zip(self.pars, pvals)
-        self.chain.config.update(pairs)
+        updated = {}
+        for idx, val in enumerate(pvals):
+            p, subkey = self.flat_to_grouped[idx]
+            if subkey is None:
+                updated[p] = val
+            else:
+                if p not in updated:
+                    updated[p] = {} if isinstance(subkey, str) else []
+                if isinstance(updated[p], dict):
+                    updated[p][subkey] = val
+                else:
+                    while len(updated[p]) <= subkey:
+                        updated[p].append(0.0)
+                    updated[p][subkey] = val
+
+        self.chain.config.update(updated)
         return
 
     def _residual(self, pvals):
@@ -118,20 +131,38 @@ def refine(self, *args, **kw):
         if not self.pars:
             return 0.0
 
-        initial = [config[p] for p in self.pars]
+        # Build flat list of initial parameters and flat_to_grouped mapping
+        initial = []
+        self.flat_to_grouped = {}
+
+        for p in self.pars:
+            val = config[p]
+            if isinstance(val, dict):
+                for k, v in val.items():
+                    initial.append(v)
+                    self.flat_to_grouped[len(initial) - 1] = (p, k)
+            elif isinstance(val, list):
+                for i, v in enumerate(val):
+                    initial.append(v)
+                    self.flat_to_grouped[len(initial) - 1] = (p, i)
+            else:
+                initial.append(val)
+                self.flat_to_grouped[len(initial) - 1] = (p, None)
+
+        # Perform least squares refinement
         sol, cov_sol, infodict, emesg, ier = leastsq(
             self.residual, initial, full_output=1
         )
         fvec = infodict["fvec"]
+
         if ier not in (1, 2, 3, 4):
-            emesg
             raise ValueError(emesg)
 
-        # Place the fit parameters in config
+        # Place the fit parameters back into config
         vals = sol
         if not hasattr(vals, "__iter__"):
             vals = [vals]
-        self.chain.config.update(zip(self.pars, vals))
+        self._update_chain(vals)
 
         return dot(fvec, fvec)
 

tests/test_morph_func.py

@@ -101,3 +101,45 @@ def test_smear_with_morph_func():
     assert np.allclose(y0, y1, atol=1e-3)  # numerical error -> 1e-4
     # verify morphed param
     assert np.allclose(smear, morphed_cfg["smear"], atol=1e-1)
+
+
+def test_squeeze_with_morph_func():
+    squeeze_init = [0, -0.001, -0.0001, 0.0001]
+    x_morph = np.linspace(0, 10, 101)
+    y_morph = 2 * np.sin(
+        x_morph + x_morph * (-0.01) - 0.0001 * x_morph**2 + 0.0002 * x_morph**3
+    )
+    expected_squeeze = [0, -0.01, -0.0001, 0.0002]
+    expected_scale = 1 / 2
+    x_target = x_morph.copy()
+    y_target = np.sin(x_target)
+    cfg = morph_default_config(scale=1.1, squeeze=squeeze_init)  # off init
+    morph_rv = morph(x_morph, y_morph, x_target, y_target, **cfg)
+    morphed_cfg = morph_rv["morphed_config"]
+    # verified they are morphable
+    x1, y1, x0, y0 = morph_rv["morph_chain"].xyallout
+    assert np.allclose(x0, x1)
+    assert np.allclose(y0, y1, atol=1e-3)  # numerical error -> 1e-4
+    # verify morphed param
+    assert np.allclose(expected_squeeze, morphed_cfg["squeeze"], atol=1e-4)
+    assert np.allclose(expected_scale, morphed_cfg["scale"], atol=1e-4)
+
+
+def test_funcy_with_morph_func():
+    def linear_function(x, y, scale, offset):
+        return (scale * x) * y + offset
+
+    x_morph = np.linspace(0, 10, 101)
+    y_morph = np.sin(x_morph)
+    x_target = x_morph.copy()
+    y_target = np.sin(x_target) * 2 * x_target + 0.4
+    cfg = morph_default_config(parameters={"scale": 1.2, "offset": 0.1})
+    cfg["function"] = linear_function
+    morph_rv = morph(x_morph, y_morph, x_target, y_target, **cfg)
+    morphed_cfg = morph_rv["morphed_config"]
+    x1, y1, x0, y0 = morph_rv["morph_chain"].xyallout
+    assert np.allclose(x0, x1)
+    assert np.allclose(y0, y1, atol=1e-6)
+    fitted_parameters = morphed_cfg["parameters"]
+    assert np.allclose(fitted_parameters["scale"], 2, atol=1e-6)
+    assert np.allclose(fitted_parameters["offset"], 0.4, atol=1e-6)