fdfitter: Improve inversion performance (Friday afternoon)

lumicks/pylake/fitting/detail/derivative_manipulation.py

@@ -21,18 +21,50 @@ def numerical_jacobian(fn, parameter_vector, dx=1e-6):
     return finite_difference_jacobian
 
 
-def inversion_functions(model_function, f_min, f_max, derivative_function, tol):
-    """This function generates two functions which allow for inverting models via optimization. These functions are used
-    by functions which require inversion of a model's dependent and independent variable"""
+def inversion_functions(model_function, f_min, f_max, derivative_function, tol, rootfinding):
+    """This function generates two functions which allow for inverting models via optimization.
+    These functions are used by functions which require inversion of a model's dependent and
+    independent variable
+
+    Parameters
+    ----------
+    model_function : callable
+        Callable that returns the dependent value when an independent value is provided.
+    f_min, f_max : float
+        Minimum and maximum value for the independent variable
+    derivative_function : callable
+        Callable that returns the derivative of the model w.r.t. independent variable.
+    tol : float
+        Tolerance
+    rootfinding : bool
+        Use Brent's method for rootfinding before going for full non-linear optimization. This
+        generally works well for monotonic functions.
+    """
+
+    def invert_brent(single_distance, _):
+        """Invert a single independent / dependent data point"""
+        return scipy.optimize.root_scalar(
+            lambda f: model_function(np.array(f)) - single_distance,
+            method="brentq",
+            bracket=(f_min, f_max),
+            rtol=tol,
+            xtol=tol,
+        ).root
 
     def fit_single(single_distance, initial_guess):
         """Invert a single independent / dependent data point"""
+        if rootfinding:
+            try:
+                return invert_brent(single_distance, initial_guess)
+            except ValueError:  # Solution outside f_min and f_max, do least_squares instead
+                pass
+
         jac = derivative_function if derivative_function else "2-point"
         single_estimate = scipy.optimize.least_squares(
             lambda f: model_function(f) - single_distance,
             initial_guess,
             jac=jac,
-            bounds=(f_min, f_max),
+            bounds=(f_min, f_max) if rootfinding else (0, np.inf),
             method="trf",
             ftol=tol,
             xtol=tol,
@@ -51,9 +83,11 @@ def manual_inversion(distances, initial_guess):
     return manual_inversion, fit_single
 
 
-def invert_function(d, initial, f_min, f_max, model_function, derivative_function=None, tol=1e-8):
-    """This function inverts a function using a least squares optimizer. For models where this is required, this is the
-    most time consuming step.
+def invert_function(
+    d, initial, f_min, f_max, model_function, derivative_function=None, tol=1e-8, rootfinding=False
+):
+    """This function inverts a function using a least squares optimizer. For models where this is
+    required, this is the most time consuming step.
 
     Parameters
     ----------
@@ -71,16 +105,27 @@ def invert_function(d, initial, f_min, f_max, model_function, derivative_functio
         model derivative with respect to the independent variable (returns an element per data point)
     tol : float
         optimization tolerances
+    rootfinding : bool
+        use rootfinding method (note that in this case, f_min and f_max are used to bracket the
+        search space, but are not enforced as hard constraints).
     """
     manual_inversion, _ = inversion_functions(
-        model_function, f_min, f_max, derivative_function, tol
+        model_function, f_min, f_max, derivative_function, tol, rootfinding
     )
 
     return manual_inversion(d, initial)
 
 
 def invert_function_interpolation(
-    d, initial, f_min, f_max, model_function, derivative_function=None, tol=1e-8, dx=1e-2
+    d,
+    initial,
+    f_min,
+    f_max,
+    model_function,
+    derivative_function=None,
+    tol=1e-8,
+    dx=1e-2,
+    rootfinding=False,
 ):
     """This function inverts a function using interpolation. For models where this is required, this is the most time
     consuming step. Specifying a sensible f_max for this method is crucial.
@@ -103,23 +148,25 @@ def invert_function_interpolation(
         desired step-size of the dependent variable
     tol : float
         optimization tolerances
+    rootfinding : bool
+        use rootfinding method (note that in this case, f_min and f_max are used to bracket the
+        search space, but are not enforced as hard constraints).
     """
     manual_inversion, fit_single = inversion_functions(
-        model_function, f_min, f_max, derivative_function, tol
+        model_function, f_min, f_max, derivative_function, tol, rootfinding
     )
     f_min_data = max([f_min, fit_single(np.min(d), initial)])
     f_max_data = min([f_max, fit_single(np.max(d), initial)])
 
     # Determine the points that lie within the range where it is reasonable to interpolate
     interpolated_idx = np.full(d.shape, False, dtype=bool)
-    f_range = np.arange(f_min_data, f_max_data, dx)
+    f_range = np.arange(f_min_data, f_max_data + dx, dx)
     if len(f_range) > 0:
         d_range = model_function(f_range)
-        d_min = np.min(d_range)
-        d_max = np.max(d_range)
+        d_min, d_max = np.min(d_range), np.max(d_range)
 
         # Interpolate for the points where interpolation is sensible
-        interpolated_idx = np.logical_and(d > d_min, d < d_max)
+        interpolated_idx = np.logical_and(d >= d_min, d <= d_max)
 
     result = np.zeros(d.shape)
     if np.sum(interpolated_idx) > 3 and len(f_range) > 3:
@@ -136,9 +183,9 @@ def invert_function_interpolation(
         result[interpolated_idx] = manual_inversion(d[interpolated_idx], initial)
 
     # Do the manual inversion for the others
-    result[np.logical_not(interpolated_idx)] = manual_inversion(
-        d[np.logical_not(interpolated_idx)], initial
-    )
+    slow_inversions = np.logical_not(interpolated_idx)
+    if np.any(slow_inversions):
+        result[slow_inversions] = manual_inversion(d[slow_inversions], initial)
 
     return result
 

lumicks/pylake/fitting/detail/model_implementation.py

@@ -725,7 +725,7 @@ def twlc_solve_force(d, Lp, Lc, St, C, g0, g1, Fc, kT=4.11):
             "Persistence length, contour length, stretch modulus and kT must be bigger than 0"
         )
 
-    f_min = 0
+    f_min = 1e-6
     f_max = (-g0 + np.sqrt(St * C)) / g1  # Above this force the model loses its validity
 
     return invert_function_interpolation(
@@ -735,6 +735,8 @@ def twlc_solve_force(d, Lp, Lc, St, C, g0, g1, Fc, kT=4.11):
         f_max,
         lambda f_trial: twlc_distance(f_trial, Lp, Lc, St, C, g0, g1, Fc, kT),
         lambda f_trial: twlc_distance_derivative(f_trial, Lp, Lc, St, C, g0, g1, Fc, kT),
+        dx=1e-2,
+        rootfinding=True,
     )
 
 

lumicks/pylake/fitting/model.py

@@ -1,5 +1,6 @@
 import uuid
 import inspect
+import warnings
 from copy import deepcopy
 from collections import OrderedDict
 
@@ -260,7 +261,14 @@ def __repr__(self):
 
         return model_info
 
-    def invert(self, independent_min=0.0, independent_max=np.inf, interpolate=False):
+    def invert(
+        self,
+        independent_min=None,
+        independent_max=None,
+        interpolate=False,
+        method="exact",
+        independent_step=1e-2,
+    ):
         """Invert this model.
 
         This operation swaps the dependent and independent parameter and should be avoided if a
@@ -274,16 +282,54 @@ def invert(self, independent_min=0.0, independent_max=np.inf, interpolate=False)
 
         Parameters
         ----------
-        independent_min : float
+        independent_min : float, optional
             Minimum value for the independent variable over which to interpolate. Only used when
             `interpolate` is set to `True`.
-        independent_max : float
+        independent_max : float, optional
             Maximum value for the independent variable over which to interpolate. Only used when
             `interpolate` is set to `True`.
-        interpolate : bool
+        interpolate : bool, optional
             Use interpolation method rather than numerical inversion.
+        method : {"exact", "interp_root", "interp_lsq"}
+            - "exact" : Invert each point separately (exact, but slow).
+            - "interp_root" : Interpolate forward model as lookup table for inversion. Appropriate range is selected by rootfinding.
+            - "interp_lsq" : Interpolate forward model as lookup table for inversion. Appropriate range is selected by least_squares (deprecated).
+              Can only be used when track is equidistantly sampled.
+
+        independent_step : float, optional
+            Interpolation step size. Default: 1e-2
         """
-        return InverseModel(self, independent_min, independent_max, interpolate)
+        if method not in ("exact", "interp_root", "interp_lsq"):
+            raise RuntimeError("Interpolation method should either be rootfinding or least_squares")
+
+        if interpolate:
+            method = "exact"
+            warnings.warn(
+                RuntimeWarning("The parameter `interpolate` is deprecated. Use method parameter.")
+            )
+
+        if method == "exact":
+            return InverseModel(
+                self, independent_min, independent_max, False, independent_step, rootfinding=False
+            )
+        elif method == "interp_root":
+            return InverseModel(
+                self,
+                1e-4 if independent_min is None else independent_min,
+                1e5 if independent_max is None else independent_max,
+                True,
+                independent_step,
+                rootfinding=True,
+            )
+        elif method == "interp_lsq":
+            return InverseModel(
+                self,
+                0.0 if independent_min is None else independent_min,
+                np.inf if independent_max is None else independent_max,
+                True,
+                independent_step,
+                rootfinding=False,
+            )
 
     def subtract_independent_offset(self):
         """
@@ -621,7 +667,15 @@ def derivative(self, independent, param_vector):
 
 
 class InverseModel(Model):
-    def __init__(self, model, independent_min=0.0, independent_max=np.inf, interpolate=False):
+    def __init__(
+        self,
+        model,
+        independent_min=0.0,
+        independent_max=np.inf,
+        interpolate=False,
+        independent_step=1e-2,
+        rootfinding=False,
+    ):
         """
         Combine two model outputs to form a new model (addition).
 
@@ -635,6 +689,10 @@ def __init__(self, model, independent_min=0.0, independent_max=np.inf, interpola
             Note that a finite maximum has to be specified if you wish to use the interpolation mode.
         interpolate : bool
             Use interpolation approximation. Default: False.
+        independent_step : float, optional
+            Interpolation step size. Default: 1e-2
+        rootfinding : bool
+            Use rootfinding rather than least squares to invert model.
 
         Raises
         ------
@@ -648,11 +706,8 @@ def __init__(self, model, independent_min=0.0, independent_max=np.inf, interpola
         self.interpolate = interpolate
         self.independent_min = independent_min
         self.independent_max = independent_max
-        if self.interpolate:
-            if not np.isfinite(independent_min) or not np.isfinite(independent_max):
-                raise ValueError(
-                    "Inversion limits have to be finite when using interpolation method."
-                )
+        self.independent_step = independent_step
+        self.rootfinding = rootfinding
 
     @property
     def dependent(self):
@@ -689,6 +744,8 @@ def _raw_call(self, independent, param_vector):
                 self.independent_max,
                 lambda f_trial: self.model._raw_call(f_trial, param_vector),
                 lambda f_trial: self.model.derivative(f_trial, param_vector),
+                dx=self.independent_step,
+                rootfinding=self.rootfinding,
             )
         else:
             return invert_function(
@@ -698,6 +755,7 @@ def _raw_call(self, independent, param_vector):
                 self.independent_max,
                 lambda f_trial: self.model._raw_call(f_trial, param_vector),  # Forward model
                 lambda f_trial: self.model.derivative(f_trial, param_vector),
+                rootfinding=self.rootfinding,
             )
 
     @property

lumicks/pylake/fitting/tests/test_fd_models.py

@@ -59,7 +59,7 @@ def test_models():
     # Check the tWLC and inverted tWLC model
     params = [5, 5, 5, 3, 2, 1, 6, 4.11]
     assert twlc_distance("tWLC").verify_jacobian(independent, params)
-    assert twlc_force("itWLC").verify_jacobian(independent, params)
+    assert twlc_force("itWLC").verify_jacobian(independent, params, rtol=1e-4)
 
     # Check whether the twistable wlc model manipulates the data order
     np.testing.assert_allclose(

lumicks/pylake/fitting/tests/test_model_composition.py

@@ -134,21 +134,14 @@ def test_subtract_independent_offset_unit(model, param, unit):
     assert model.defaults[param].unit == unit
 
 
-def test_interpolation_inversion():
-    m = ewlc_odijk_distance("Nucleosome").invert(independent_max=120.0, interpolate=True)
+@pytest.mark.parametrize("method", ["exact", "interp_lsq", "interp_root"])
+def test_interpolation_inversion(method):
+    m = ewlc_odijk_distance("Nucleosome").invert(independent_max=120.0, method=method)
     parvec = [5.77336105517341, 7.014180463612673, 1500.0000064812095, 4.11]
     result = np.array([0.17843862, 0.18101283, 0.18364313, 0.18633117, 0.18907864])
     np.testing.assert_allclose(m._raw_call(np.arange(10, 250, 50) / 1000, parvec), result)
 
 
-@pytest.mark.parametrize("param", [{"independent_max": np.inf}, {"independent_min": -np.inf}])
-def test_interpolation_invalid_range(param):
-    with pytest.raises(
-        ValueError, match="Inversion limits have to be finite when using interpolation method"
-    ):
-        ewlc_odijk_distance("Nucleosome").invert(**param, interpolate=True)
-
-
 def test_uuids():
     m1, m2 = (Model(name, lambda x: x, dependent="x") for name in ("M1", "M2"))
     m3 = CompositeModel(m1, m2)