From c0e3f8b3530ee6a22a3cfda745cca9358814cea2 Mon Sep 17 00:00:00 2001
From: Tom de Geus <tom@geus.me>
Date: Mon, 7 Mar 2022 13:48:47 +0100
Subject: [PATCH] Adding `fit_log`

---
 GooseMPL/__init__.py | 30 +++++++++++++++++++++++++++---
 test/main.py         | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/GooseMPL/__init__.py b/GooseMPL/__init__.py
index b9354fc..2302973 100644
--- a/GooseMPL/__init__.py
+++ b/GooseMPL/__init__.py
@@ -1319,7 +1319,7 @@ def fit_powerlaw(
         with ``x`` replaced with the specified string.
 
     :param extrapolate:
-        Plot the powerlaw on the full range of ``axis.get_xlim()``.
+        Plot the function on the full range of ``axis.get_xlim()``.
         Instead of ``True``, one can specify plot options for the extrapolated line, e.g.
         ``..., extrapolate=dict(ls="--", c="r"), ...``.
 
@@ -1471,7 +1471,7 @@ def fit_exp(
         with ``x`` replaced with the specified string.
 
     :param extrapolate:
-        Plot the powerlaw on the full range of ``axis.get_xlim()``.
+        Plot the function on the full range of ``axis.get_xlim()``.
         Instead of ``True'', one can specify plot options for the extrapolated line, e.g.
         ``..., extrapolate=dict(ls="--", c="r"), ...``.
 
@@ -1570,6 +1570,30 @@ def fit_exp(
     return (prefactor, exponent, details)
 
 
+def fit_log(
+    xdata: ArrayLike,
+    ydata: ArrayLike,
+    **kwargs,
+) -> (float, float, dict):
+    r"""
+    Fit a logarithm :math:`y = a + b \ln x`.
+    See documentation of :py:func:`fit_linear`.
+    """
+
+    xdata = np.array(xdata)
+    ydata = np.array(ydata)
+
+    i = xdata > 0
+    logx = np.log(xdata[i])
+    y = ydata[i]
+
+    j = np.isnan(logx)
+    logx = logx[~j]
+    y = y[~j]
+
+    return fit_linear(logx, y, **kwargs)
+
+
 def fit_linear(
     xdata: ArrayLike,
     ydata: ArrayLike,
@@ -1605,7 +1629,7 @@ def fit_linear(
         with ``x`` replaced with the specified string.
 
     :param extrapolate:
-        Plot the powerlaw on the full range of ``axis.get_xlim()``.
+        Plot the function on the full range of ``axis.get_xlim()``.
         Instead of ``True``, one can specify plot options for the extrapolated line, e.g.
         ``..., extrapolate=dict(ls="--", c="r"), ...``.
 
diff --git a/test/main.py b/test/main.py
index 30110fe..fd3b5f4 100644
--- a/test/main.py
+++ b/test/main.py
@@ -190,6 +190,44 @@ def test_exponent(self):
         self.assertTrue(np.isclose(exponent, 3.4))
 
 
+class Test_fit_log(unittest.TestCase):
+    """
+    Fit an logarithmic function.
+    """
+
+    def test_prefactor_exponent(self):
+
+        x = np.linspace(0, 1, 1000)[1:]
+        y = 1.2 + 3.4 * np.log(x)
+        offset, prefactor, _ = gplt.fit_log(x, y)
+        self.assertTrue(np.isclose(offset, 1.2))
+        self.assertTrue(np.isclose(prefactor, 3.4))
+
+    def test_prefactor_negative_prefactor(self):
+
+        x = np.linspace(0, 1, 1000)[1:]
+        y = 1.2 - 3.4 * np.log(x)
+        offset, prefactor, _ = gplt.fit_log(x, y)
+        self.assertTrue(np.isclose(offset, 1.2))
+        self.assertTrue(np.isclose(prefactor, -3.4))
+
+    def test_prefactor(self):
+
+        x = np.linspace(0, 1, 1000)[1:]
+        y = 1.2 + 3.4 * np.log(x)
+        offset, prefactor, _ = gplt.fit_log(x, y, prefactor=3.4)
+        self.assertTrue(np.isclose(offset, 1.2))
+        self.assertTrue(np.isclose(prefactor, 3.4))
+
+    def test_exponent(self):
+
+        x = np.linspace(0, 1, 1000)[1:]
+        y = 1.2 + 3.4 * np.log(x)
+        offset, prefactor, _ = gplt.fit_log(x, y, offset=1.2)
+        self.assertTrue(np.isclose(offset, 1.2))
+        self.assertTrue(np.isclose(prefactor, 3.4))
+
+
 class Test_fit_linear(unittest.TestCase):
     """
     Fit a linear.