Repository: matthewwardrop/formulaic

Issue:
When working with time series data, is there a way to use lagged variables?
Let's say we have time series data and two variables in a Pandas dataframe. Let's say we want to run a regression of some target variable $y_t$ on $x_t$ and $y_{t-1}$. Is there a way to do something like the following code?

`LHS, RHS = model_matrix("y ~ x + lag(y, 1)")`

I took a look on the documentation and couldn't find it.

Hints from issue comments:
@matthewwardrop: Hi @rgriva ! We don't have inbuilt support for this, but maybe we should. In any case, you can do:

```
def lag(x, shift):
    return x.shift(shift)

LHS, RHS = model_matrix("y ~ x  + lag(y, 1)", df)
```

Apologies for the delay in my response!@matthewwardrop: I think I'll go ahead and add this to the standard set of transforms for the next release, since this is a pretty common use-case.@matthewwardrop: This has now been added for the next release of Formulaic, which will be 1.1.0, and should be released somewhat soon. Thanks for reaching out @rgriva !

Solution patch:
diff --git a/docsite/docs/guides/grammar.md b/docsite/docs/guides/grammar.md
index eee53726..710f1e43 100644
--- a/docsite/docs/guides/grammar.md
+++ b/docsite/docs/guides/grammar.md
@@ -62,6 +62,7 @@ that have *not* been implemented by `formulaic` are explicitly noted also.
 | `center(...)` | Shift column data so mean is zero. | ✓ | ✓ | ✗ |
 | `scale(...)` | Shift column so mean is zero and variance is 1. | ✓ | ✓[^7] | ✓ |
 | `standardize(...)` | Alias of `scale`. | ✓[^8] | ✓ | ✗ |
+| `lag(...[, <k>])` | Generate lagging or leading columns (useful for datasets collected at regular intervals). | ✓ | ✗ | ✓ |
 | `poly(...)` | Generates a polynomial basis, allowing non-linear fits. | ✓ | ✗ | ✓ |
 | `bs(...)` | Generates a B-Spline basis, allowing non-linear fits. | ✓ | ✓ | ✓ |
 | `cr(...)` | Generates a natural cubic spline basis, allowing non-linear fits. | ✗ | ✓ | ✓ |
diff --git a/docsite/docs/guides/transforms.ipynb b/docsite/docs/guides/transforms.ipynb
index ba1b2ff0..5def3b7e 100644
--- a/docsite/docs/guides/transforms.ipynb
+++ b/docsite/docs/guides/transforms.ipynb
@@ -48,6 +48,7 @@
     "    * **exp10**: `numpy.exp10`.\n",
     "    * **exp2**: `numpy.exp2`.\n",
     "    * **I**: Identity/null transform (alternative to `{<expr>}` syntax).\n",
+    "    * **lag**: Generate lagging or leading columns (useful for datasets collected at regular intervals).\n",
     "* Stateful transforms (documented below):\n",
     "    * **bs**: Basis spline coding.\n",
     "    * **center**: Subtraction of the mean.\n",
diff --git a/formulaic/transforms/__init__.py b/formulaic/transforms/__init__.py
index 85502bef..66cfaa13 100644
--- a/formulaic/transforms/__init__.py
+++ b/formulaic/transforms/__init__.py
@@ -6,6 +6,7 @@
 from .contrasts import C, ContrastsRegistry, encode_contrasts
 from .hashed import hashed
 from .identity import identity
+from .lag import lag
 from .patsy_compat import PATSY_COMPAT_TRANSFORMS
 from .poly import poly
 from .scale import center, scale
@@ -16,6 +17,7 @@
     "C",
     "encode_contrasts",
     "ContrastsRegistry",
+    "lag",
     "poly",
     "center",
     "scale",
@@ -35,6 +37,7 @@
     # Bespoke transforms
     "bs": basis_spline,
     "center": center,
+    "lag": lag,
     "poly": poly,
     "scale": scale,
     "C": C,
diff --git a/formulaic/transforms/lag.py b/formulaic/transforms/lag.py
new file mode 100644
index 00000000..5aae7cbe
--- /dev/null
+++ b/formulaic/transforms/lag.py
@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import functools
+from typing import Any, TYPE_CHECKING
+
+import numpy
+import pandas
+
+
+@functools.singledispatch
+def lag(data: Any, k: int = 1) -> Any:
+    """
+    Shift the data indices forward or backward by the specified number of
+    intervals.
+
+    Args:
+        data: The array-like dataset to be lagged along its first axis.
+        k: The number of intervals to shift back the observations by. For "lead"
+            shifts, use negative integers.
+
+    Notes:
+      * For this transform to make sense, `data` must represent a time-series
+        collected at regularly spaced intervals.
+      * `k` is chosen as the offset parameter for compatibility with R.
+    """
+    raise NotImplementedError(
+        f"No implementation of `shift` for data of type {type(data)}."
+    )
+
+
+@lag.register
+def _(data: pandas.Series, offset: int = 1) -> pandas.Series:
+    return data.shift(offset)
+
+
+@lag.register
+def _(data: numpy.ndarray, offset: int = 1) -> numpy.ndarray:
+    if data.shape == () or offset == 0:
+        return data
+
+    data = numpy.copy(data).astype(float)
+
+    if offset > 0:
+        data[offset:] = data[:-offset]
+        data[:offset] = numpy.nan
+    else:
+        data[:offset] = data[-offset:]
+        data[offset:] = numpy.nan
+
+    return data


Test patch:
diff --git a/tests/transforms/test_lag.py b/tests/transforms/test_lag.py
new file mode 100644
index 00000000..5c15ac56
--- /dev/null
+++ b/tests/transforms/test_lag.py
@@ -0,0 +1,38 @@
+import numpy
+import pandas
+import pytest
+
+from formulaic.transforms import lag
+
+TEST_CASES = [
+    ([1, 2, 3, 4], None, [numpy.nan, 1, 2, 3]),
+    ([1, 2, 3, 4], 1, [numpy.nan, 1, 2, 3]),
+    ([1, 2, 3, 4], 2, [numpy.nan, numpy.nan, 1, 2]),
+    ([1, 2, 3, 4], 0, [1, 2, 3, 4]),
+    ([1, 2, 3, 4], -1, [2, 3, 4, numpy.nan]),
+]
+
+
+@pytest.mark.parametrize("data, k, expected", TEST_CASES)
+def test_lag_series(data, k, expected):
+    assert numpy.allclose(
+        lag(pandas.Series(data), k).values
+        if k is not None
+        else lag(pandas.Series(data)).values,
+        expected,
+        equal_nan=True,
+    )
+
+
+@pytest.mark.parametrize("data, k, expected", TEST_CASES)
+def test_lag_array(data, k, expected):
+    assert numpy.allclose(
+        lag(numpy.array(data), k) if k is not None else lag(numpy.array(data)),
+        expected,
+        equal_nan=True,
+    )
+
+
+def test_lag_invalid_type():
+    with pytest.raises(NotImplementedError):
+        lag(None)