From 0c1537711e25cfa7ede5d201dac20312e6454878 Mon Sep 17 00:00:00 2001
From: Joep Vanlier <joep.vanlier@gmail.com>
Date: Mon, 7 Oct 2024 18:48:22 +0200
Subject: [PATCH] calibration: expose surface drag models

---
 lumicks/pylake/__init__.py                    |  1 +
 .../force_calibration/calibration_models.py   | 32 +++++++++++++++++++
 .../force_calibration/detail/drag_models.py   |  6 ++--
 .../force_calibration/tests/test_hydro.py     | 11 +++++++
 4 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/lumicks/pylake/__init__.py b/lumicks/pylake/__init__.py
index 39e27c908..97c846872 100644
--- a/lumicks/pylake/__init__.py
+++ b/lumicks/pylake/__init__.py
@@ -34,6 +34,7 @@
     density_of_water,
     viscosity_of_water,
     coupling_correction_2d,
+    surface_drag_correction,
 )
 from .force_calibration.power_spectrum_calibration import (
     fit_power_spectrum,
diff --git a/lumicks/pylake/force_calibration/calibration_models.py b/lumicks/pylake/force_calibration/calibration_models.py
index 0a7799cf3..f8045ff40 100644
--- a/lumicks/pylake/force_calibration/calibration_models.py
+++ b/lumicks/pylake/force_calibration/calibration_models.py
@@ -180,6 +180,38 @@ def viscosity_of_water(temperature, molarity_nacl=None, pressure=None):
         return _poly((temperature + 273.15) / 300, bi, ai) * 1e-6
 
 
+def surface_drag_correction(distance_to_surface, bead_diameter, axial=False):
+    """Calculate a correction factor for the drag coefficient due to a nearby surface [1]_.
+
+    Parameters
+    ----------
+    distance_to_surface : array_like | float
+        Distance from the center of the bead to the surface
+    bead_diameter : float
+        Bead diameter
+    axial : bool
+        Compute the correction factor for axial drag
+
+    References
+    ----------
+    .. [1] Schäffer, E., Nørrelykke, S. F., & Howard, J. "Surface forces and drag coefficients of
+           microspheres near a plane surface measured with optical tweezers." Langmuir, 23(7), 3654-3665
+           (2007).
+    """
+    distance_to_surface = np.asarray(distance_to_surface)
+
+    if np.any(distance_to_surface <= 0.5 * bead_diameter):
+        raise ValueError(
+            f"Bead is inside the surface. you specified a distance of {distance_to_surface}, while"
+            f"the bead radius is {bead_diameter / 2}"
+        )
+
+    if axial:
+        return brenner_axial(distance_to_surface, bead_diameter / 2)
+    else:
+        return faxen_factor(distance_to_surface, bead_diameter / 2)
+
+
 def coupling_correction_2d(dx, dy, bead_diameter, is_y_oscillation=False, allow_rotation=True):
     """Calculates the coupling correction factor for a 2D problem.
 
diff --git a/lumicks/pylake/force_calibration/detail/drag_models.py b/lumicks/pylake/force_calibration/detail/drag_models.py
index f859fe8bb..818a722d2 100644
--- a/lumicks/pylake/force_calibration/detail/drag_models.py
+++ b/lumicks/pylake/force_calibration/detail/drag_models.py
@@ -14,7 +14,7 @@ def faxen_factor(distance_to_surface_m, radius_m):
 
     Parameters
     ----------
-    distance_to_surface_m : float
+    distance_to_surface_m : array_like | float
         Distance from the center of the bead to the surface [m]
     radius_m : float
         Radius of the bead [m]
@@ -31,7 +31,7 @@ def faxen_factor(distance_to_surface_m, radius_m):
 
 
 def brenner_axial(distance_to_surface_m, radius_m):
-    """Brenner factor for lateral drag coefficient.
+    """Brenner factor for axial drag coefficient.
 
     This factor provides a correction to the drag force for a nearby wall.
 
@@ -41,7 +41,7 @@ def brenner_axial(distance_to_surface_m, radius_m):
 
     Parameters
     ----------
-    distance_to_surface_m : float
+    distance_to_surface_m : array_like | float
         Distance from the center of the bead to the surface [m]
     radius_m : float
         Radius of the bead [m]
diff --git a/lumicks/pylake/force_calibration/tests/test_hydro.py b/lumicks/pylake/force_calibration/tests/test_hydro.py
index 2bb67fba4..cc93dd8c3 100644
--- a/lumicks/pylake/force_calibration/tests/test_hydro.py
+++ b/lumicks/pylake/force_calibration/tests/test_hydro.py
@@ -3,6 +3,7 @@
 from lumicks.pylake.force_calibration.calibration_models import (
     ActiveCalibrationModel,
     PassiveCalibrationModel,
+    surface_drag_correction,
 )
 from lumicks.pylake.force_calibration.detail.drag_models import *
 from lumicks.pylake.force_calibration.detail.power_models import g_diode
@@ -492,6 +493,11 @@ def test_faxen(bead_radius, result):
     np.testing.assert_allclose(
         faxen_factor(bead_radius + np.arange(40, 250, 80) * 1e-9, bead_radius), result
     )
+    np.testing.assert_allclose(
+        surface_drag_correction(bead_radius + np.arange(40, 250, 80) * 1e-9, bead_radius),
+        result,
+        False,
+    )
 
 
 @pytest.mark.parametrize(
@@ -505,6 +511,11 @@ def test_brenner(bead_radius, result):
     np.testing.assert_allclose(
         brenner_axial(bead_radius + np.arange(40, 250, 80) * 1e-9, bead_radius), result
     )
+    np.testing.assert_allclose(
+        surface_drag_correction(bead_radius + np.arange(40, 250, 80) * 1e-9, bead_radius),
+        result,
+        True,
+    )
 
 
 def test_near_surface_consistency():