From 833c2e47364498c7902668e282344d3a56237735 Mon Sep 17 00:00:00 2001
From: Clemens Prescher <clemens.prescher@gmail.com>
Date: Sun, 14 Jul 2024 08:51:50 +0200
Subject: [PATCH] s0 can now be set manually for extrapolation to zero

---
 glassure/calc.py          |  6 +++++-
 glassure/configuration.py |  1 +
 tests/test_calc.py        | 10 ++++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/glassure/calc.py b/glassure/calc.py
index fdb2d43..dba1c88 100644
--- a/glassure/calc.py
+++ b/glassure/calc.py
@@ -108,7 +108,11 @@ def process_input(input: Input) -> Pattern:
     sq = calculate_sq(norm, f_squared_mean, f_mean_squared)
 
     # extrapolation
-    s0 = calculate_s0(composition)
+    if config.transform.extrapolation.s0 is not None:
+        s0 = config.transform.extrapolation.s0
+    else:
+        s0 = calculate_s0(composition)
+        
     extrapolation = transform.extrapolation
     match extrapolation.method:
         case ExtrapolationMethod.STEP:
diff --git a/glassure/configuration.py b/glassure/configuration.py
index 15f83a6..d6cffde 100644
--- a/glassure/configuration.py
+++ b/glassure/configuration.py
@@ -49,6 +49,7 @@ class OptimizeConfig:
 @dataclass
 class ExtrapolationConfig:
     method: ExtrapolationMethod = ExtrapolationMethod.STEP
+    s0: Optional[float] = field(default=None)
     overlap: float = 0.2
     replace: bool = False
 
diff --git a/tests/test_calc.py b/tests/test_calc.py
index 20ed638..bd74fad 100644
--- a/tests/test_calc.py
+++ b/tests/test_calc.py
@@ -97,6 +97,16 @@ def test_process_input_poly_extrapolation():
 
     assert not np.array_equal(res.sq.y, res_poly.sq.y)
 
+def test_process_extrapolation_with_s0():
+    input = prepare_input()
+    res = process_input(input)
+
+    input.config.transform.extrapolation.method = "linear"
+    input.config.transform.extrapolation.s0 = 0.1
+    res_s0 = process_input(input)
+
+    assert not np.array_equal(res.sq.y, res_s0.sq.y)
+
 
 def test_process_input_kn_correction():
     input = prepare_input()