abtoolkit 1.0.0__tar.gz

abtoolkit-1.0.0/LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 Nikita Altukhov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

abtoolkit-1.0.0/PKG-INFO

@@ -0,0 +1,113 @@
+Metadata-Version: 2.1
+Name: abtoolkit
+Version: 1.0.0
+Summary: Package with tools for AB testing
+Author-email: Nikita Altukhov <altuxov.nikita@gmail.com>
+Project-URL: Homepage, https://github.com/nikitosl/abtoolkit
+Project-URL: Issues, https://github.com/nikitosl/abtoolkit/issues
+Keywords: ab_test,cuped,did,ttest
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pandas>=2.2.1
+Requires-Dist: numpy>=1.26.4
+Requires-Dist: statsmodels>=0.14.1
+Requires-Dist: scipy>=1.12.0
+Requires-Dist: linearmodels>=5.4
+Requires-Dist: tqdm>=4.66.2
+
+# ABToolkit
+Set of tools for AA and AB tests, sample size estimation, confidence intervals estimation. 
+For continuous and discrete variables.
+
+## Install using pip:
+```pip install abtoolkit```
+
+## Continuous variables analysis
+#### Sample size estimation:
+```
+from abtoolkit.continuous.utils import calculate_sample_size_by_mde
+calculate_sample_size_by_mde(
+    std=variable.std(),
+    alpha=alpha_level, 
+    power=power, 
+    mde=mde
+)
+```
+
+#### AA and AB tests simulation:
+Using ```abtoolkit.continuous.simulation.StatTestsSimulation``` class you can simulate and check different stat-test, 
+compare them in terms of stat test power to choose the best test for your data. As result of simulation for each 
+stat test you will get the 1-st Type error estimation with confidence interval, 2-nd Type error estimation with 
+confidence interval and plot of p-value distribution for different tests.
+
+```
+from abtoolkit.continuous.simulation import StatTestsSimulation
+simulation = StatTestsSimulation(
+    control,
+    test,
+    stattests_list=["ttest", "regression_test", "cuped_ttest", "did_regression_test", "additional_vars_regression_test"],
+    experiments_num=experiments_num,
+    sample_size=sample_size,
+    mde=mde,
+    alpha_level=alpha_level,
+
+    control_previous_values=control_previous_value,
+    test_previous_values=test_previous_value,
+    control_cuped_covariant=control_previous_value,
+    test_cuped_covariant=test_previous_value,
+    control_additional_vars=[control_previous_value],
+    test_additional_vars=[test_previous_value],
+)
+simulation.run()  # Run simulation
+simulation.print_results()  # Print results of simulation
+simulation.plot_p_values()  # Plot p-values distribution
+```
+Output:
+![output-plot.png](static%2Foutput-plot.png)
+![p-value-plot.png](static%2Fp-value-plot.png)
+
+Full example of usage you can find in ```examples/continuous_var_analysis.py``` script.
+
+#### Next stat tests implemented for treatment effect estimation:
+- ***T-Test*** - estimates treatment effect by comparing variables between test and control groups.
+- ***Difference T-Test*** - estimates treatment effect by comparing difference between actual and previous values 
+of variables in test and control groups.
+- ***Regression Test*** - estimates treatment effect using linear regression by tested predicting variable. 
+Fact of treatment represented in model as binary flag (treated or not). Weight for this flag show significant 
+of treatment impact.
+```y = bias + w * treated```
+- ***Regression Difference-in-Difference Test*** - estimates treatment effect using linear regression by predicting
+difference between test and control groups whist represented as difference between current variable value and 
+previous period variable value (two differences). Weight for treated and current variable values shows 
+significant of treatment. ```y = bias + w0 * treated + w1 * after + w2 * treated * after```
+- ***CUPED*** - estimates treatment effect by comparing variables between test and control groups and uses covariant 
+to reduce variance and speedup test. ```y = y - Q * covariant```, where ```Q = cov(y, covariant) / var(covariant)```. 
+Cuped variable has same mean value (unbiased), but smaller variance, that speedup test.
+- ***Regression with Additional Variables*** - estimates treatment effect using linear regression by predicting 
+tested variable with additional variables, which describe part of main variable variance and speedup test. 
+Fact of treatment represented in model as binary flag (treated or not). Weight for this flag show significant 
+of treatment impact.
+```y = bias + w0 * treated + w1 * additional_variable1 + w2 * additional_variable2 + ...```
+
+
+## Discrete variables analysis
+#### Sample size estimation:
+```
+from abtoolkit.discrete.utils import estimate_ci_binomial
+estimate_ci_binomial(
+    p, 
+    sample_size, 
+    alpha=0.05
+)
+```
+#### AA and AB tests simulation:
+To Be Done
+#### Next stat tests implemented for treatment effect estimation:
+To Be Done
+
+---
+You can find examples of toolkit usage in [examples/](https://github.com/nikitosl/abtoolkit/tree/master/examples) directory.

abtoolkit-1.0.0/README.md

@@ -0,0 +1,92 @@
+# ABToolkit
+Set of tools for AA and AB tests, sample size estimation, confidence intervals estimation. 
+For continuous and discrete variables.
+
+## Install using pip:
+```pip install abtoolkit```
+
+## Continuous variables analysis
+#### Sample size estimation:
+```
+from abtoolkit.continuous.utils import calculate_sample_size_by_mde
+calculate_sample_size_by_mde(
+    std=variable.std(),
+    alpha=alpha_level, 
+    power=power, 
+    mde=mde
+)
+```
+
+#### AA and AB tests simulation:
+Using ```abtoolkit.continuous.simulation.StatTestsSimulation``` class you can simulate and check different stat-test, 
+compare them in terms of stat test power to choose the best test for your data. As result of simulation for each 
+stat test you will get the 1-st Type error estimation with confidence interval, 2-nd Type error estimation with 
+confidence interval and plot of p-value distribution for different tests.
+
+```
+from abtoolkit.continuous.simulation import StatTestsSimulation
+simulation = StatTestsSimulation(
+    control,
+    test,
+    stattests_list=["ttest", "regression_test", "cuped_ttest", "did_regression_test", "additional_vars_regression_test"],
+    experiments_num=experiments_num,
+    sample_size=sample_size,
+    mde=mde,
+    alpha_level=alpha_level,
+
+    control_previous_values=control_previous_value,
+    test_previous_values=test_previous_value,
+    control_cuped_covariant=control_previous_value,
+    test_cuped_covariant=test_previous_value,
+    control_additional_vars=[control_previous_value],
+    test_additional_vars=[test_previous_value],
+)
+simulation.run()  # Run simulation
+simulation.print_results()  # Print results of simulation
+simulation.plot_p_values()  # Plot p-values distribution
+```
+Output:
+![output-plot.png](static%2Foutput-plot.png)
+![p-value-plot.png](static%2Fp-value-plot.png)
+
+Full example of usage you can find in ```examples/continuous_var_analysis.py``` script.
+
+#### Next stat tests implemented for treatment effect estimation:
+- ***T-Test*** - estimates treatment effect by comparing variables between test and control groups.
+- ***Difference T-Test*** - estimates treatment effect by comparing difference between actual and previous values 
+of variables in test and control groups.
+- ***Regression Test*** - estimates treatment effect using linear regression by tested predicting variable. 
+Fact of treatment represented in model as binary flag (treated or not). Weight for this flag show significant 
+of treatment impact.
+```y = bias + w * treated```
+- ***Regression Difference-in-Difference Test*** - estimates treatment effect using linear regression by predicting
+difference between test and control groups whist represented as difference between current variable value and 
+previous period variable value (two differences). Weight for treated and current variable values shows 
+significant of treatment. ```y = bias + w0 * treated + w1 * after + w2 * treated * after```
+- ***CUPED*** - estimates treatment effect by comparing variables between test and control groups and uses covariant 
+to reduce variance and speedup test. ```y = y - Q * covariant```, where ```Q = cov(y, covariant) / var(covariant)```. 
+Cuped variable has same mean value (unbiased), but smaller variance, that speedup test.
+- ***Regression with Additional Variables*** - estimates treatment effect using linear regression by predicting 
+tested variable with additional variables, which describe part of main variable variance and speedup test. 
+Fact of treatment represented in model as binary flag (treated or not). Weight for this flag show significant 
+of treatment impact.
+```y = bias + w0 * treated + w1 * additional_variable1 + w2 * additional_variable2 + ...```
+
+
+## Discrete variables analysis
+#### Sample size estimation:
+```
+from abtoolkit.discrete.utils import estimate_ci_binomial
+estimate_ci_binomial(
+    p, 
+    sample_size, 
+    alpha=0.05
+)
+```
+#### AA and AB tests simulation:
+To Be Done
+#### Next stat tests implemented for treatment effect estimation:
+To Be Done
+
+---
+You can find examples of toolkit usage in [examples/](https://github.com/nikitosl/abtoolkit/tree/master/examples) directory.

abtoolkit-1.0.0/abtoolkit/continuous/simulation.py

@@ -0,0 +1,269 @@
+from typing import List
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+
+from abtoolkit.continuous.stattests import additional_vars_regression_test
+from abtoolkit.continuous.stattests import cuped_ttest
+from abtoolkit.continuous.stattests import did_regression_test
+from abtoolkit.continuous.stattests import regression_test
+from abtoolkit.continuous.stattests import ttest
+from abtoolkit.discrete.utils import estimate_ci_binomial
+
+
+class StatTestsSimulation:
+    def __init__(
+            self,
+            control: pd.Series,
+            test: pd.Series,
+            stattests_list: List[str],
+            sample_size: int,
+            experiments_num: int,
+            mde: float,
+            alpha_level: float = 0.05,
+            power: float = 0.8,
+            control_previous_values: pd.Series = None,
+            test_previous_values: pd.Series = None,
+            control_cuped_covariant: pd.Series = None,
+            test_cuped_covariant: pd.Series = None,
+            control_additional_vars: List[pd.Series] = None,
+            test_additional_vars: List[pd.Series] = None,
+    ):
+        """
+        Simulates AA and AB tests for given stat-tests. Prints result (alpha and power) for each test
+        and builds plot for p-value distributions.
+
+        :param control: control variable
+        :param test: test variable
+        :param stattests_list: list of stat-tests for estimation
+        :param sample_size: number of examples to sample from variables in each iteration
+        :param experiments_num: number of experiments to perform for each stat-test
+        :param mde: minimal detectable effect, used to perform AB test (add to test variable)
+        :param alpha_level: test alpha-level
+        :param power: test power
+        :param control_previous_values: previous values of control variable used to reduce variance and speedup
+        test in difference-in-difference test
+        :param test_previous_values: previous values of test variable used to reduce variance and speedup test
+        in difference-in-difference test
+        :param control_cuped_covariant: covariant for control group variable used to reduce variance and speedup test
+        in cuped test
+        :param test_cuped_covariant: covariant for test group variable used to reduce variance and speedup test
+        in cuped test
+        :param control_additional_vars: list of additional variables for control group variable used to
+        reduce variance and speedup test in 'regression_with_additional_variables' test
+        :param test_additional_vars: list of additional variables for test group variable used to
+        reduce variance and speedup test in 'regression_with_additional_variables' test
+        """
+
+        self.control = control
+        self.test = test
+
+        self.stattests_list = stattests_list
+        self.experiments_num = experiments_num
+        self.sample_size = sample_size
+        self.mde = mde
+        self.alpha_level = alpha_level
+        self.power = power
+
+        self.stattests_func_map = {
+            "ttest": self.simulate_ttest,
+            "diff_ttest": self.simulate_difference_ttest,
+            "cuped_ttest": self.simulate_cuped,
+            "regression_test": self.simulate_reg,
+            "did_regression_test": self.simulate_reg_did,
+            "additional_vars_regression_test": self.simulate_reg_add,
+        }
+        self.info = {}
+
+        # Optional
+        self.control_previous_values = control_previous_values
+        self.test_previous_values = test_previous_values
+        self.control_cuped_covariant = control_cuped_covariant
+        self.test_cuped_covariant = test_cuped_covariant
+        self.control_additional_vars = control_additional_vars
+        self.test_additional_vars = test_additional_vars
+
+    def plot_p_values(self):
+        """
+        Plot p-values distribution for each test
+        :return: None
+        """
+        if len(self.info) == 0:
+            return
+
+        X = np.linspace(0, 1, 1000)
+        for test, test_info in self.info.items():
+            ab_pvalues = np.array(test_info["ab_pvalues"])
+            Y = [np.mean(ab_pvalues < x) for x in X]
+            plt.plot(X, Y, label=test)
+
+        plt.plot([self.alpha_level, self.alpha_level], [0, 1], '--k', alpha=0.8)
+        plt.plot([0, 1], [self.power, self.power], '--k', alpha=0.8)
+        plt.title('P-Value Distribution for AB Simulation', size=12)
+        plt.xlabel('p-value', size=10)
+        plt.legend(fontsize=10)
+        plt.grid()
+        plt.show()
+
+    def print_results(self):
+        """
+        Print simulation results for each test (alpha and power + confidence intervals)
+        :return: None
+        """
+        for test in self.info:
+            a, p = self.info[test]['alpha'], self.info[test]['power']
+            aci1, aci2 = round(self.info[test]['alpha_ci'][0], 4), round(self.info[test]['alpha_ci'][1], 4)
+            pci1, pci2 = round(self.info[test]['power_ci'][0], 4), round(self.info[test]['power_ci'][1], 4)
+
+            if (aci1 > self.alpha_level) or (self.power > pci2):
+                print('\033[91m' + f"'{test}'; alpha={a} ci[{aci1}; {aci2}], power={p} [{pci1}; {pci2}]" + '\033[0m')
+            else:
+                print('\033[92m' + f"'{test}'; alpha={a} ci[{aci1}; {aci2}], power={p} [{pci1}; {pci2}]" + '\033[0m')
+
+    def run(self):
+        """
+        Simulate all tests from 'self.stattests_list' by given data and save information to 'info' dictionary
+        :return:
+        """
+        self.info = {}
+        for stattest in self.stattests_list:
+            self.simulate_test_by_name(stattest)
+        return self.info
+
+    def simulate_test_by_name(self, test_name: str):
+        """
+        Simulate AA and AB test and save results to 'info' dictionary
+        :param test_name: name of test for simulation (ttest | cuped_ttest | regression_test | did_regression_test
+                                                        | additional_vars_regression_test)
+        :return: None
+        """
+
+        assert test_name in self.stattests_func_map, f"Given test_name {test_name} not found"
+        stattest_func = self.stattests_func_map[test_name]
+
+        test_success_no_effect_cnt = 0
+        test_pvalues_no_effect = []
+        test_success_effect_cnt = 0
+        test_pvalues_effect = []
+
+        for _ in tqdm(range(self.experiments_num), desc=f"Simulation test '{test_name}'"):
+            p_value = stattest_func(mde=0)
+            test_pvalues_no_effect.append(p_value)
+            if p_value < self.alpha_level:
+                test_success_no_effect_cnt += 1
+
+            p_value = stattest_func(mde=self.mde)
+            test_pvalues_effect.append(p_value)
+            if p_value < self.alpha_level:
+                test_success_effect_cnt += 1
+    
+        alpha = test_success_no_effect_cnt / self.experiments_num
+        power = test_success_effect_cnt / self.experiments_num
+
+        alpha_ci = estimate_ci_binomial(alpha, self.experiments_num, alpha=0.05)
+        power_ci = estimate_ci_binomial(power, self.experiments_num, alpha=0.05)
+
+        if test_name in self.info:
+            del self.info[test_name]
+        self.info[test_name] = {
+            "alpha": alpha,
+            "alpha_ci": alpha_ci,
+            "power": power,
+            "power_ci": power_ci,
+            "aa_pvalues": test_pvalues_no_effect,
+            "ab_pvalues": test_pvalues_effect,
+        }
+
+    def simulate_ttest(self, mde: float) -> float:
+        """
+        Simulate ttest
+        :param mde: minimal detectable effect, to sum with test variable
+        :return: p_value
+        """
+        control_sample = self.control.sample(self.sample_size, replace=True)
+        test_sample = self.test.sample(self.sample_size, replace=True)
+        test_sample += mde
+    
+        return ttest(control_sample, test_sample)
+
+    def simulate_difference_ttest(self, mde: float) -> float:
+        """
+        Simulate ttest for difference between actual variable value and previous period variable value
+        :param mde: minimal detectable effect, to sum with test variable
+        :return: p_value
+        """
+        control_index_sample = self.control.index[np.random.randint(0, len(self.control), size=self.sample_size)]
+        test_index_sample = self.test.index[np.random.randint(0, len(self.test), size=self.sample_size)]
+
+        control_sample = self.control.loc[control_index_sample]
+        control_pre_sample = self.control_previous_values.loc[control_index_sample]
+        test_sample = self.test.loc[test_index_sample]
+        test_pre_sample = self.test_previous_values.loc[test_index_sample]
+        test_sample += mde
+
+        return cuped_ttest(control_sample, control_pre_sample, test_sample, test_pre_sample)
+
+    def simulate_cuped(self, mde: float) -> float:
+        """
+        Simulate CUPED ttest
+        :param mde: minimal detectable effect, to sum with test variable
+        :return: p_value
+        """
+        control_index_sample = self.control.index[np.random.randint(0, len(self.control), size=self.sample_size)]
+        test_index_sample = self.test.index[np.random.randint(0, len(self.test), size=self.sample_size)]
+    
+        control_sample = self.control.loc[control_index_sample]
+        control_covariant_sample = self.control_cuped_covariant.loc[control_index_sample]
+        test_sample = self.test.loc[test_index_sample]
+        test_covariant_sample = self.test_cuped_covariant.loc[test_index_sample]
+        test_sample += mde
+    
+        return cuped_ttest(control_sample, control_covariant_sample, test_sample, test_covariant_sample)
+
+    def simulate_reg(self, mde: float) -> float:
+        """
+        Simulate test using regression
+        :param mde: minimal detectable effect, to sum with test variable
+        :return: p_value
+        """
+        control_sample = self.control.sample(self.sample_size, replace=True)
+        test_sample = self.test.sample(self.sample_size, replace=True)
+        test_sample += mde
+
+        return regression_test(control_sample, test_sample)
+
+    def simulate_reg_did(self, mde: float) -> float:
+        """
+        Simulate test using regression with difference-in-difference technique
+        :param mde: minimal detectable effect, to sum with test variable
+        :return: p_value
+        """
+        control_index_sample = self.control.index[np.random.randint(0, len(self.control), size=self.sample_size)]
+        test_index_sample = self.test.index[np.random.randint(0, len(self.test), size=self.sample_size)]
+
+        control_sample = self.control.loc[control_index_sample]
+        control_previous_sample = self.control_previous_values.loc[control_index_sample]
+        test_sample = self.test.loc[test_index_sample]
+        test_previous_sample = self.test_previous_values.loc[test_index_sample]
+        test_sample += mde
+
+        return did_regression_test(control_sample, control_previous_sample, test_sample, test_previous_sample)
+
+    def simulate_reg_add(self, mde: float) -> float:
+        """
+        Simulate test using regression with additional variables
+        :param mde: minimal detectable effect, to sum with test variable
+        :return: p_value
+        """
+        control_index_sample = self.control.index[np.random.randint(0, len(self.control), size=self.sample_size)]
+        test_index_sample = self.test.index[np.random.randint(0, len(self.test), size=self.sample_size)]
+
+        control_sample = self.control.loc[control_index_sample]
+        control_add_samples = [a.loc[control_index_sample] for a in self.control_additional_vars]
+        test_sample = self.test.loc[test_index_sample]
+        test_add_samples = [a.loc[test_index_sample] for a in self.test_additional_vars]
+        test_sample += mde
+
+        return additional_vars_regression_test(control_sample, control_add_samples, test_sample, test_add_samples)

abtoolkit-1.0.0/abtoolkit/continuous/stattests.py

@@ -0,0 +1,192 @@
+from typing import List
+
+import linearmodels as lm
+import numpy as np
+import pandas as pd
+from scipy.stats import ttest_ind
+
+
+def ttest(
+        control: pd.Series,
+        test: pd.Series,
+) -> float:
+    """
+    Simple two-side t-test
+    :param control: pd.Series for control sample
+    :param test: pd.Series for test sample
+    :return: p-value
+    """
+
+    return ttest_ind(control, test, alternative="less").pvalue
+
+
+def difference_ttest(
+        control: pd.Series,
+        control_pre: pd.Series,
+        test: pd.Series,
+        test_pre: pd.Series,
+) -> float:
+    """
+    Estimation treatment effect using ttest and CUPED to increase test's power
+    :param control: pd.Series, control sample
+    :param control_pre: pd.Series, control previous period value
+    :param test: pd.Series, test sample
+    :param test_pre: pd.Series, test previous period value
+    :return: p-value
+    """
+    control = control - control_pre
+    test = test - test_pre
+
+    return ttest(control, test)
+
+
+def cuped_ttest(
+        control: pd.Series,
+        control_covariant: pd.Series,
+        test: pd.Series,
+        test_covariant: pd.Series,
+) -> float:
+    """
+    Estimation treatment effect using ttest and CUPED to increase test's power
+    :param control: pd.Series, control sample
+    :param control_covariant: pd.Series, control sample covariant
+    :param test: pd.Series, test sample
+    :param test_covariant: pd.Series, test sample covariant
+    :return: p-value
+    """
+
+    full_value = pd.concat([
+        control.rename("value"),
+        test.rename("value"),
+    ], axis=0)
+
+    full_covariant = pd.concat([
+        control_covariant.rename("covariant"),
+        test_covariant.rename("covariant"),
+    ], axis=0)
+
+    cov = np.cov(full_covariant, full_value)[0, 1]
+    var = full_covariant.var()
+    theta = cov / var
+
+    cuped_test = test - theta * test_covariant
+    cuped_control = control - theta * control_covariant
+
+    return ttest(cuped_control, cuped_test)
+
+
+
+def regression_test(
+        control: pd.Series,
+        test: pd.Series,
+) -> float:
+    """
+    Treatment effect estimation using linear regression
+    :param control: pd.Series with index [entity, dt], where dt could be int of datetime. Control sample
+    :param test: pd.Series with index [entity, dt], where dt could be int of datetime. Test sample
+    :return: p-value
+    """
+    df = pd.concat([
+        control.rename("value").to_frame().assign(treated=0),
+        test.rename("value").to_frame().assign(treated=1),
+    ], axis=0)
+    df["bias"] = 1
+
+    if not isinstance(df.index, pd.MultiIndex):
+        df["index1"] = 0
+        df["index2"] = 1
+        df = df.set_index(["index1", "index2"])
+
+    mod = lm.PanelOLS.from_formula("value ~ bias + treated", data=df)
+    result = mod.fit()
+    return result.pvalues["treated"]
+
+
+def did_regression_test(
+        control: pd.Series,
+        control_pre: pd.Series,
+        test: pd.Series,
+        test_pre: pd.Series,
+) -> float:
+    """
+    Difference-in-Difference treatment effect estimation using linear regression.
+    Calculates difference between current and last values in test and control groups and then
+    calculates difference between differences to increase test power
+    :param control_pre: pd.Series with index [entity, dt], where dt could be int of datetime.
+    Control sample before treatment
+    :param control: pd.Series with index [entity, dt], where dt could be int of datetime.
+    Control sample after treatment
+    :param test_pre: pd.Series with index [entity, dt], where dt could be int of datetime. Test sample before treatment
+    :param test: pd.Series with index [entity, dt], where dt could be int of datetime. Test sample after treatment
+    :return: p-value
+    """
+    df = pd.concat([
+        control_pre.rename("value").to_frame().assign(treated=0).assign(after=0),
+        control.rename("value").to_frame().assign(treated=0).assign(after=1),
+        test_pre.rename("value").to_frame().assign(treated=1).assign(after=0),
+        test.rename("value").to_frame().assign(treated=1).assign(after=1),
+    ], axis=0)
+
+    df["bias"] = 1
+
+    if not isinstance(df.index, pd.MultiIndex):
+        df["index1"] = 0
+        df["index2"] = 1
+        df = df.set_index(["index1", "index2"])
+
+    mod = lm.PanelOLS.from_formula("value ~ bias + + after + treated + treated*after", data=df)
+    result = mod.fit()
+    return result.pvalues["treated:after"]
+
+
+def additional_vars_regression_test(
+        control: pd.Series,
+        control_additional_vars: List[pd.Series],
+        test: pd.Series,
+        test_additional_vars: List[pd.Series],
+) -> float:
+    """
+    Treatment effect estimation using additional variables in linear regression. Additional
+    variables should reduce deviation of target variable and increase test power
+    :param control: pd.Series with index [entity, dt], where dt could be int of datetime.
+    Control sample
+    :param control_additional_vars: List of pd.Series with index [entity, dt], where dt could be int of datetime.
+    Additional variables which can describe some deviation of tested variable
+    :param test: pd.Series with index [entity, dt], where dt could be int of datetime.
+    Test sample
+    :param test_additional_vars: List of pd.Series with index [entity, dt], where dt could be int of datetime.
+    Additional variables which can describe some deviation of tested variable
+    :return: p-value
+    """
+
+    assert len(test_additional_vars) > 0, "No additional vars for 'additional_vars_regression_test' test given"
+
+    additional_vars_names_test = [v.name for v in test_additional_vars]
+    additional_vars_names_control = [v.name for v in control_additional_vars]
+    assert set(additional_vars_names_test) == set(additional_vars_names_control), \
+        (f"Lists of control and test additional vars should the same. "
+         f"Got {set(additional_vars_names_test)} vars for test "
+         f"and {set(additional_vars_names_control)} vars for control")
+
+    control_df = pd.concat([control.rename("value").to_frame()] + control_additional_vars, axis=1)
+    test_df = pd.concat([test.rename("value").to_frame()] + test_additional_vars, axis=1)
+
+    control_df.index = test_df.index
+    df = pd.concat([
+        control_df.assign(treated=0),
+        test_df.assign(treated=1),
+    ], axis=0)
+
+    df["bias"] = 1
+
+    if not isinstance(df.index, pd.MultiIndex):
+        df["index1"] = 0
+        df["index2"] = 1
+        df = df.set_index(["index1", "index2"])
+
+    additional_vars_formula = " + ".join(map(str, additional_vars_names_test))
+
+    formula = f"value ~ bias + treated + {additional_vars_formula}"
+    mod = lm.PanelOLS.from_formula(formula, data=df)
+    result = mod.fit()
+    return result.pvalues["treated"]