From 3da1f6cfb40b7afb336763cb565471835c7102b0 Mon Sep 17 00:00:00 2001
From: Niklas Selke <n.selke@fz-juelich.de>
Date: Thu, 15 Sep 2022 09:56:25 +0200
Subject: [PATCH] Added OLS and quantile regression.
---
toarstats/trends/interface.py | 21 ++++++++++++++++++++-
toarstats/trends/ols.py | 22 ++++++++++++++++++++++
toarstats/trends/quant_reg.py | 23 +++++++++++++++++++++++
3 files changed, 65 insertions(+), 1 deletion(-)
create mode 100644 toarstats/trends/ols.py
create mode 100644 toarstats/trends/quant_reg.py
diff --git a/toarstats/trends/interface.py b/toarstats/trends/interface.py
index d2787cd..b6b928a 100644
--- a/toarstats/trends/interface.py
+++ b/toarstats/trends/interface.py
@@ -4,7 +4,8 @@ This module contains the following function:
calculate_trend - calculate the trend using the requested method
"""
-import statsmodels.formula.api as smf
+from toarstats.trends.ols import ols
+from toarstats.trends.quant_reg import quant_reg
def calculate_trend(method, data, formula="value ~ datetime", quantiles=None):
@@ -22,4 +23,22 @@ def calculate_trend(method, data, formula="value ~ datetime", quantiles=None):
:param quantiles: a single quantile or a list of quantiles to
calculate, these must be between 0 and 1; only
needed when ``method="quant"``
+
+ :raises ValueError: raised if the ``method`` parameter is not
+ recognized
+
+ :return: The result of the fit or a list of fit results if
+ ``method="quant"`` and multiple quantiles are given
"""
+ if method == "OLS":
+ fit = ols(data, formula)
+ elif method == "quant":
+ try:
+ fit = [quant_reg(data, formula, quantile)
+ for quantile in quantiles]
+ except TypeError:
+ fit = quant_reg(data, formula, quantiles)
+ else:
+ raise ValueError(f"{method} is not recognized, must be 'OLS' or"
+ " 'quant'")
+ return fit
diff --git a/toarstats/trends/ols.py b/toarstats/trends/ols.py
new file mode 100644
index 0000000..255741d
--- /dev/null
+++ b/toarstats/trends/ols.py
@@ -0,0 +1,22 @@
+"""Ordinary least squares (OLS) linear regression calculation.
+
+This module contains the following function:
+ols - calculate the OLS linear regression
+"""
+
+import statsmodels.formula.api as smf
+
+
+def ols(data, formula):
+ """Calculate the OLS linear regression.
+
+ :param data: data containing a list of date time values and
+ associated parameter values on which to calculate the
+ trend
+ :param formula: the formula specifying the model
+
+ :return: The result of the fit
+ """
+ model = smf.ols(formula, data)
+ fit = model.fit()
+ return fit
diff --git a/toarstats/trends/quant_reg.py b/toarstats/trends/quant_reg.py
new file mode 100644
index 0000000..0335977
--- /dev/null
+++ b/toarstats/trends/quant_reg.py
@@ -0,0 +1,23 @@
+"""Quantile regression calculation.
+
+This module contains the following function:
+quant_reg - calculate the OLS linear regression
+"""
+
+import statsmodels.formula.api as smf
+
+
+def quant_reg(data, formula, quantile):
+ """Calculate the quantile regression.
+
+ :param data: data containing a list of date time values and
+ associated parameter values on which to calculate the
+ trend
+ :param formula: the formula specifying the model
+ :param quantile: a single quantile, must be between 0 and 1
+
+ :return: The result of the fit
+ """
+ model = smf.quantreg(formula, data)
+ fit = model.fit(q=quantile)
+ return fit
--
GitLab