reformulate lambda function

mesmer/stats/_power_transformer.py

@@ -29,7 +29,7 @@ def lambda_function(xi_0, xi_1, local_yearly_T):
     lambdas : ndarray of float of shape (n_years,)
         The parameters of the power transformation for each gridcell and month
     """
-    return 2 / (1 + xi_0 * np.exp(local_yearly_T * xi_1))
+    return 2 / (1 + np.exp(xi_0) * np.exp(local_yearly_T * xi_1))
 
 
 def _yeo_johnson_transform_np(data, lambdas):
@@ -153,14 +153,14 @@ def _neg_log_likelihood(coeffs):
 
         return -loglikelihood
 
-    bounds = np.array([[0, np.inf], [-0.1, 0.1]])
-    first_guess = np.array([1, 0])
+    # bounds = np.array([[0, np.inf], [-0.1, 0.1]])
+    first_guess = np.array([0, 0])
 
     xi_0, xi_1 = minimize(
         _neg_log_likelihood,
         x0=first_guess,
-        bounds=bounds,
-        method="Nelder-Mead",
+        # bounds=bounds,
+        method="BFGS",
     ).x
 
     return xi_0, xi_1

tests/unit/test_power_transformer.py

@@ -15,24 +15,48 @@
 
 
 @pytest.mark.parametrize(
-    "coeffs, t, expected",
+    "coeffs, T, expected",
     [
-        ([1, 0.1], -np.inf, 2.0),
-        ([1, 0.1], np.inf, 0.0),
-        ([1, -0.1], -np.inf, 0.0),
-        ([1, -0.1], np.inf, 2.0),
-        ([0, 0], 1, 2),
-        ([0, 1], 1, 2),
-        ([1, 0], 1, 1),
-        ([2, 0], 1, 2 / 3),
-        ([1, 1], np.log(9), 2 / 10),
+        (
+            [1, 0.1],
+            np.inf,
+            0.0,
+        ),  # coeffs for monotonic increase function goes to zero for positive infinity
+        (
+            [1, 0.1],
+            -np.inf,
+            2.0,
+        ),  # coeffs for monotonic increase function goes to 2 for negative infinity
+        (
+            [1, -0.1],
+            np.inf,
+            2.0,
+        ),  # coeffs for monotonic decrease function goes to 2 for positive infinity
+        (
+            [1, -0.1],
+            -np.inf,
+            0.0,
+        ),  # coeffs for monotonic decrease function goes to zero for negative infinity
+        ([0, 0], 1, 1.0),  # for zero zero function is a constant at 1
+        ([0, 0], np.pi, 1.0),  # for zero zero ALL values are 1
+        (
+            [0, 1],
+            np.log(9),
+            2 / 10,
+        ),  # for 0, 1 the function simplifies to 2 / (1+ exp(x))
     ],
 )
-def test_lambda_function(coeffs, t, expected):
+def test_lambda_function(coeffs, T, expected):
 
-    result = mesmer.stats.lambda_function(coeffs[0], coeffs[1], t)
+    result = mesmer.stats.lambda_function(coeffs[0], coeffs[1], T)
     np.testing.assert_allclose(result, expected)
 
+    def old_lambda_function(coeffs, T):
+        return 2 / (1 + coeffs[0] * np.exp(coeffs[1] * T))
+
+    result_old = old_lambda_function([np.exp(coeffs[0]), coeffs[1]], T)
+    np.testing.assert_allclose(result, result_old)
+
 
 def test_yeo_johnson_optimize_lambda_np_normal():
     # with enough random normal data points the fit should be close to 1 and 0
@@ -44,12 +68,12 @@ def test_yeo_johnson_optimize_lambda_np_normal():
 
     result = _yeo_johnson_optimize_lambda_np(monthly_residuals, yearly_T)
     # to test viability
-    expected = np.array([1, 0])
+    expected = np.array([0, 0])
     np.testing.assert_allclose(result, expected, atol=1e-2)
 
     # to test numerical stability
-    expected_exact = np.array([9.976913e-01, -1.998520e-05])
-    np.testing.assert_allclose(result, expected_exact, atol=1e-7)
+    expected_exact = np.array([-0.001162, -0.001162])  # [[-0.001099, -0.001224]])
+    np.testing.assert_allclose(result, expected_exact, atol=1e-6)
 
 
 @pytest.mark.parametrize(
@@ -76,6 +100,9 @@ def test_yeo_johnson_optimize_lambda_np(skew, bounds):
     transformed = _yeo_johnson_transform_np(local_monthly_residuals, lmbda)
 
     assert (lmbda >= bounds[0]).all() & (lmbda <= bounds[1]).all()
+    assert np.abs(sp.stats.skew(transformed)) <= np.abs(
+        sp.stats.skew(local_monthly_residuals)
+    )
     np.testing.assert_allclose(sp.stats.skew(transformed), 0, atol=0.1)
 
 
@@ -203,6 +230,7 @@ def skewed_data_2D(n_timesteps=30, n_lat=3, n_lon=2):
 
     ts_array = np.empty((n_cells, n_timesteps))
     rng = np.random.default_rng(0)
+    np.random.seed(0)  # for scipy
 
     # create random data with a skew
     skew = rng.uniform(-5, 5, size=(n_cells, 1))