[ENH] Change effects api, and use BaseObject from scikit-base

docs/examples/custom-trend.ipynb

@@ -10,8 +10,7 @@
     "tags": []
    },
    "source": [
-    "# Use-case: forecasting diffusion of innovations\n",
-    "\n",
+    "# Use-case: forecasting product adoption\n",
     "\n",
     "**Diffusion of innovations** is a theory that seeks to explain how, why, and at what rate new ideas and technology spread through cultures. This theory was formulated by E.M. Rogers in 1962 and is often used to understand the adoption or spread of new products and technologies among different groups of people. \n",
     "\n",
@@ -51,11 +50,11 @@
     "\n",
     "This repository had an initial explosion of stars during the first ~10 days, which we will ignore since the daily granularity is not enough to capture the initial growth (hourly might work). After that, the number of starts grew by following a bell-shaped curve, which we will try to model. This curve might be related to the popularity of deep learning itself.\n",
     "\n",
-    "```{info}\n",
-    "\n",
-    "This dataset was obtained from https://github.com/emanuelef/daily-stars-explorer. \n",
+    "<div class=\"admonition note\">\n",
+    "    <p class=\"admonition-title\">Note</p>\n",
+    "    <p>This dataset was obtained from https://github.com/emanuelef/daily-stars-explorer. </p>\n",
+    "</div>\n",
     "\n",
-    "```\n",
     "\n"
    ]
   },
@@ -1157,7 +1156,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.19"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

extension_templates/effect.py

@@ -0,0 +1,96 @@
+"""Extension template for creating a new effect."""
+
+from typing import Any, Dict
+
+import jax.numpy as jnp
+import pandas as pd
+
+from prophetverse.effects.base import BaseEffect, Stage
+from prophetverse.utils.frame_to_array import series_to_tensor_or_array
+
+
+class MyEffectName(BaseEffect):
+    """Base class for effects."""
+
+    _tags = {
+        # Supports multivariate data? Can this
+        # Effect be used with Multiariate prophet?
+        "supports_multivariate": False,
+        # If no columns are found, should
+        # _predict be skipped?
+        "skip_predict_if_no_match": True,
+    }
+
+    def __init__(self, param1: Any, param2: Any):
+        self.param1 = param1
+        self.param2 = param2
+
+    def _fit(self, X: pd.DataFrame, scale: float = 1.0):
+        """Customize the initialization of the effect.
+
+        This method is called by the `fit()` method and can be overridden by
+        subclasses to provide additional initialization logic.
+
+        Parameters
+        ----------
+        X : pd.DataFrame
+            The DataFrame to initialize the effect.
+        """
+        # Do something with X, scale, and other parameters
+        pass
+
+    def _transform(
+        self, X: pd.DataFrame, stage: Stage = Stage.TRAIN
+    ) -> Dict[str, jnp.ndarray]:
+        """Prepare the input data in a dict of jax arrays.
+
+        This method is called by the `fit()` method and can be overridden
+        by subclasses to provide additional data preparation logic.
+
+        Parameters
+        ----------
+        X : pd.DataFrame
+            The input DataFrame containing the exogenous variables for the training
+            time indexes, if passed during fit, or for the forecasting time indexes, if
+            passed during predict.
+
+        stage : Stage, optional
+            The stage of the effect, by default Stage.TRAIN. This can be used to
+            differentiate between training and prediction stages and apply different
+            transformations accordingly.
+
+        Returns
+        -------
+        Dict[str, jnp.ndarray]
+            A dictionary containing the data needed for the effect. The keys of the
+            dictionary should be the names of the arguments of the `apply` method, and
+            the values should be the corresponding data as jnp.ndarray.
+        """
+        # Do something with X
+        if stage == "train":
+            array = series_to_tensor_or_array(X)
+        else:
+            # something else
+            pass
+        return {"data": array}
+
+    def _predict(self, trend: jnp.ndarray, **kwargs) -> jnp.ndarray:
+        """Apply the effect.
+
+        This method is called by the `apply()` method and must be overridden by
+        subclasses to provide the actual effect computation logic.
+
+        Parameters
+        ----------
+        trend : jnp.ndarray
+            An array containing the trend values.
+
+        kwargs: dict
+            Additional keyword arguments that may be needed to compute the effect.
+
+        Returns
+        -------
+        jnp.ndarray
+            The effect values.
+        """
+        raise NotImplementedError("Subclasses must implement _predict()")

mkdocs.yml

@@ -10,6 +10,7 @@ nav:
   - Examples:
       - Univariate Time Series: examples/univariate.ipynb
       - Hierarchical Time Series: examples/hierarchical.ipynb
+      - Custom Effects: examples/custom-effect.ipynb
       - Count Time Series: examples/count-data.ipynb
       - Custom Trend: examples/custom-trend.ipynb
   - Reference:
@@ -41,6 +42,9 @@ plugins:
       use_directory_urls: false
 
 markdown_extensions:
+  - admonition
+  - pymdownx.details
+  - pymdownx.superfences
   - pymdownx.arithmatex:
         generic: true
 

pyproject.toml

@@ -12,7 +12,7 @@ sktime = ">=0.30.0"
 numpyro = ">=0.13.2"
 optax = ">=0.2"
 graphviz = "^0.20.3"
-
+scikit-base = "^0.8.1"
 
 ipykernel = { version = ">=6.26.0,<7.0.0", optional = true }
 pytest = { version = ">=8.0.0,<9.0.0", optional = true }
@@ -33,6 +33,7 @@ mypy = { version = ">=1.10.0,<2.0.0", optional = true }
 pylint = { version = ">=3.2.2,<4.0.0", optional = true }
 mkdocstrings = { version = "^0.25.1", optional = true }
 
+
 [tool.poetry.extras]
 dev = [
     "ipykernel",

src/prophetverse/effects/__init__.py

@@ -1,17 +1,17 @@
 """Effects that define relationships between variables and the target."""
 
-from .base import AbstractEffect
-from .effect_apply import additive_effect, matrix_multiplication, multiplicative_effect
+from .base import BaseEffect
+from .fourier import LinearFourierSeasonality
 from .hill import HillEffect
+from .lift_experiment import LiftExperimentLikelihood
 from .linear import LinearEffect
 from .log import LogEffect
 
 __all__ = [
-    "AbstractEffect",
-    "additive_effect",
-    "multiplicative_effect",
-    "matrix_multiplication",
+    "BaseEffect",
     "HillEffect",
     "LinearEffect",
     "LogEffect",
+    "LiftExperimentLikelihood",
+    "LinearFourierSeasonality",
 ]