Repository: matthewwardrop/formulaic

Issue:
Add support for multi-stage formulas.
In some of my work I am interested in exploring two-stage least-square regression on sparse data, and thus in making Formulaic able to handle it nicely.

My plan is to allow formulas of form:
`y ~ a + [b + c ~ z1 + z2] | a + [e + f ~ z1 + z2] | d + [b + c ~ z1 + z2] | d + [e + f ~ z1 + z2]`
In my proposed grammar, this would also be equivalent to:
`y ~ (a|d) + [b + c | e + f ~ z1 + z2]`
Using multipart syntax in the rhs of nested formulas would be forbidden.

The API for accessing the various pieces of this Formula is as yet not fully fleshed out, and naming has not been properly considered, but would be something like:
```py
f = Formula('y ~ (a|d) + [b + c | e + f ~ z1 + z2]')
f.formula_for(rhs_part=0, stage=0)  # b + c ~ z1 + z2
f.formula_for(rhs_part=0, stage=1)  # y ~ a + b + c
f.formula_for(rhs_part=1, stage=0)  # e + f ~ z1 + z2
f.formula_for(rhs_part=1, stage=1)  # y ~ a + e + f
f.formula_for(rhs_part=0) # y ~ a + [b + c ~ z1 + z2]

f = Formula('y ~ x + z')
f.formula_for() # y ~ x + z
```

On a multipart formula like this one, calls to `get_model_matrix` will need to specify the part and stage for which the model matrix should be generated. If there is only one part or stage, this will not be necessary. Formulaic explicitly will not attempt to do any modeling with this, and will expect users of the library to do any memoisation that is required for two-stage least-squares to work when pumping new data sets through a pre-trained model.

I'm especially keen to know what @bashtage thinks about this, given that this is something he has explored a lot more in linearmodels.

Hints from issue comments:
@bashtage: What is the intention of the first formula? What is exogenous and what is endogenous?  Clearly the Z are instruments. 
.@matthewwardrop: Returning to this after several years :sweat: .

Multi-part formulas are already implemented as of `v0.3.0`: `y ~ a | b | c` does the right thing.

@bashtage : If I were to take this further, I'd look to implement something like: `y ~ 1 + x1 + x2 + x3  + [ x4 + x5 ~ z1 + z2 + z3]`, exactly as you have done [here](https://bashtage.github.io/linearmodels/iv/examples/using-formulas.html). The results would be made available on the `Structured` instance as something like:

```
.lhs
    y
.rhs
    1 + x1 + x2 + x3 + IV[x4] + IV[x5]
    .iv_x4:
        z1 + z2 + z3
    .iv_x5:
        z1 + z2 + z3
```

This is within reach of the parser now, but I'd love your take on this (given that you have much more experience in this space).@bashtage: An advanced syntax would be great.  I have a few current uses.

1. IV like you have above.
2. Absorbing regression where high dimensional fixed effects are absorbed.  Something like `y ~ x + [eff1 + eff2 + eff3]` where `eff#` are categorical variables usually that are then encoded to sparse arrays.
3. Systems equations.  I currently use a dictionary.  These models have multiple equations, something like `y1 ~ x + z, y2 ~ x + w`.  Not sure if something like this woudl make sense to have as a syntax. @matthewwardrop: Nice. I don't yet know how much it makes sense to *always* have these advanced operators in place (versus having a family of parsers that extend some common set), but I'll definitely be working toward making the parser *capable* of generating formulae for these kinds of situations.

For further clarity:
On 2. Absorbing regression is just your usual fixed-effects regression, right? Where you demean the data based on a set of covariates prior to modelling, perhaps using another regression? What would you want output in that case? Something like:
```
.lhs
    y_residuals
    .fixed_effects
         eff1 + eff2 + eff3
.rhs
    x
```

On 3. Would a `Structured` instance of a tuple of formulas work? That could be implemented trivially today (either in formulaic or downstream by adding the `,` operator):
```
[0]
    .rhs
        y1
    .lhs
        x + z
[1]
    ...
```@bashtage: I haven't really through about it.  I could imagine that formulas could be nested.  For example

```
y ~ 1 + x + [w ~ z]
```

could be something like

```
.lhs
   y
.rhs
   1 + x + [w ~ z]
```

and when you access `.rhs` it would be [1{Term}, x{Term}, [w ~ z]{Formula}] so that one could handle nested formulas with some recusions, e.g.

``` 
for term_or_fmla in formula.rhs.terms:
    if isinstance(term_or_fmla , Term):
        """Do something"""
    else:
        """Handle nested formula probably using a recursion"""
```

Maybe too complicted.
@GuiMarthe: Just to add to this style of syntax, mlogit uses something similar for multinomial choice models. Not saying it should be implemented here, but there is another use case for the `|` syntax. 
In that literature, y ~ x | z | w, is the notation used for the anatomy of utility functions and translates to
```
choice ~ alternative vars. with generic coefficients |
                individual vars. with specific coefficients |
                alternative vars. with specific coefficients
```@matthewwardrop: @GuiMarthe This is actually already implemented in Formulaic (leaving the interpretation to the calling library).

![image](https://github.com/matthewwardrop/formulaic/assets/124910/1ad9d3ca-e9c6-4442-8b71-991c04e1aa75)

The wrapping library would then just need to validate that the formula has the expected structure (it could also, if desired, disable the intercept additions in the formula parser).


Solution patch:
diff --git a/docsite/docs/guides/grammar.md b/docsite/docs/guides/grammar.md
index 710f1e43..fc227457 100644
--- a/docsite/docs/guides/grammar.md
+++ b/docsite/docs/guides/grammar.md
@@ -42,9 +42,10 @@ unless otherwise indicated.
 | `+` | 1 | Returns the current term unmodified (not very useful). | ✓ | ✓ | ✓ |
 | `-` | 1 | Negates a term (only implemented for 0, in which case it is replaced with `1`). | ✓ | ✓ | ✓ |
 |-----|
-| `|` | 2 | Splits a formula into multiple parts, allowing the simultaneous generation of multiple model matrices. When on the right-hand-side of the `~` operator, all parts will attract an additional intercept term by default. | ✓ | ✗ | ✓[^6] |
+| `\|` | 2 | Splits a formula into multiple parts, allowing the simultaneous generation of multiple model matrices. When on the right-hand-side of the `~` operator, all parts will attract an additional intercept term by default. | ✓ | ✗ | ✓[^6] |
 |-----|
 | `~` | 1,2 | Separates the target features from the input features. If absent, it is assumed that we are considering only the the input features. Unless otherwise indicated, it is assumed that the input features implicitly include an intercept. | ✓ | ✓ | ✓ |
+| `[ . ~ . ]` | 2 | [Experimental] Multi stage formula notation, which is useful in (e.g.) IV contexts. Requires the `MULTISTAGE` feature flag to be passed to the parser. | ✓ | ✗ | ✗ |
 
 
 ## Transforms
diff --git a/formulaic/model_spec.py b/formulaic/model_spec.py
index d1c7e2a4..adaeb493 100644
--- a/formulaic/model_spec.py
+++ b/formulaic/model_spec.py
@@ -24,8 +24,8 @@
 from .materializers import ClusterBy, FormulaMaterializer, NAAction
 
 if TYPE_CHECKING:  # pragma: no cover
-    from .transforms.contrasts import ContrastsState
     from .model_matrix import ModelMatrices, ModelMatrix
+    from .transforms.contrasts import ContrastsState
 
 # Cached property was introduced in Python 3.8 (we currently support 3.7)
 try:
diff --git a/formulaic/parser/parser.py b/formulaic/parser/parser.py
index 830f8797..e85a7f1f 100644
--- a/formulaic/parser/parser.py
+++ b/formulaic/parser/parser.py
@@ -35,9 +35,14 @@ class DefaultParserFeatureFlag(Flag):
     Feature flags to restrict the flexibility of the formula parser.
     """
 
-    NONE = 0
     TWOSIDED = auto()
     MULTIPART = auto()
+    MULTISTAGE = auto()
+
+    # Convenience flags
+    NONE = 0
+    DEFAULT = TWOSIDED | MULTIPART
+    ALL = TWOSIDED | MULTIPART | MULTISTAGE
 
     @classmethod
     def from_spec(
@@ -79,9 +84,7 @@ class DefaultFormulaParser(FormulaParser):
         default_factory=lambda: DefaultOperatorResolver()  # pylint: disable=unnecessary-lambda
     )
     include_intercept: bool = True
-    feature_flags: DefaultParserFeatureFlag = (
-        DefaultParserFeatureFlag.TWOSIDED | DefaultParserFeatureFlag.MULTIPART
-    )
+    feature_flags: DefaultParserFeatureFlag = DefaultParserFeatureFlag.DEFAULT
 
     def __post_init__(self) -> None:
         if isinstance(self.feature_flags, set):
@@ -132,13 +135,34 @@ def get_tokens(self, formula: str) -> Iterable[Token]:
                     join_operator="+",
                 )
             )
-            rhs_index = (
-                max(
-                    (i for i, token in enumerate(tokens) if token.token.endswith("~")),
-                    default=-1,
-                )
-                + 1
-            )
+
+            def find_rhs_index(tokens: List[Token]) -> int:
+                """
+                Find the top-level index of the tilde operator starting the
+                right hand side of the formula (or -1 if not found).
+                """
+                from .algos.tokens_to_ast import CONTEXT_CLOSERS, CONTEXT_OPENERS
+
+                context = []
+                for index, token in enumerate(tokens):
+                    if token.kind is Token.Kind.CONTEXT:
+                        if token.token in CONTEXT_OPENERS:
+                            context.append(token.token)
+                            continue
+                        else:
+                            if (
+                                not context
+                                or context[-1] != CONTEXT_CLOSERS[token.token]
+                            ):
+                                return -1  # pragma: no cover ; should not happen
+                            context.pop()
+                    if context:
+                        continue
+                    if token.token == "~":  # noqa: S105
+                        return index
+                return -1
+
+            rhs_index = find_rhs_index(tokens) + 1
             tokens = [
                 *(
                     tokens[:rhs_index]
@@ -238,9 +262,7 @@ class DefaultOperatorResolver(OperatorResolver):
             internally).
     """
 
-    feature_flags: DefaultParserFeatureFlag = (
-        DefaultParserFeatureFlag.TWOSIDED | DefaultParserFeatureFlag.MULTIPART
-    )
+    feature_flags: DefaultParserFeatureFlag = DefaultParserFeatureFlag.DEFAULT
 
     def __post_init__(self) -> None:
         if isinstance(self.feature_flags, set):
@@ -292,6 +314,25 @@ def power(arg: OrderedSet[Term], power: OrderedSet[Term]) -> OrderedSet[Term]:
                 for term in itertools.product(*[arg] * int(power_term.factors[0].expr))
             )
 
+        def multistage_formula(
+            lhs: OrderedSet[Term], rhs: OrderedSet[Term]
+        ) -> Structured[OrderedSet[Term]]:
+            def get_terms(terms: OrderedSet[Term]) -> List[Term]:
+                return [
+                    Term(
+                        factors=[Factor(str(t) + "_hat", eval_method="lookup")],
+                        origin=t,
+                    )
+                    for t in terms
+                ]
+
+            if isinstance(lhs, Structured):
+                raise NotImplementedError(
+                    "Nested multistage formulas do not support structured lhs."
+                )
+
+            return Structured(get_terms(lhs), deps=(Structured(lhs=lhs, rhs=rhs),))
+
         return [
             Operator(
                 "~",
@@ -313,6 +354,16 @@ def power(arg: OrderedSet[Term], power: OrderedSet[Term]) -> OrderedSet[Term]:
                 accepts_context=lambda context: len(context) == 0,
                 structural=True,
             ),
+            Operator(
+                "~",
+                arity=2,
+                precedence=-100,
+                associativity=None,
+                to_terms=multistage_formula,
+                accepts_context=lambda context: bool(context) and context[-1] == "[",
+                structural=True,
+                disabled=DefaultParserFeatureFlag.MULTISTAGE not in self.feature_flags,
+            ),
             Operator(
                 "|",
                 arity=2,
diff --git a/formulaic/parser/types/factor.py b/formulaic/parser/types/factor.py
index d8af6947..d6cad25c 100644
--- a/formulaic/parser/types/factor.py
+++ b/formulaic/parser/types/factor.py
@@ -99,4 +99,6 @@ def to_terms(self) -> OrderedSet[Term]:
         return OrderedSet((Term([self]),))
 
     def __repr__(self) -> str:
+        if ":" in self.expr:
+            return f"`{self.expr}`"
         return self.expr
diff --git a/formulaic/parser/types/term.py b/formulaic/parser/types/term.py
index 158bb233..7157689d 100644
--- a/formulaic/parser/types/term.py
+++ b/formulaic/parser/types/term.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Iterable
+import re
+from typing import TYPE_CHECKING, Any, Iterable, Optional
 
 if TYPE_CHECKING:
     from .factor import Factor  # pragma: no cover
@@ -15,10 +16,15 @@ class Term:
 
     Attributes:
         factors: The set of factors to be multiplied to form the term.
+        origin: If this `Term` has been derived from another `Term`, for example
+            in subformulae, a reference to the original term.
     """
 
-    def __init__(self, factors: Iterable["Factor"]):
+    FACTOR_MATCHER = re.compile(r"(?:^|(?<=:))(`?)(?P<factor>[^`]+?)\1(?=:|$)")
+
+    def __init__(self, factors: Iterable["Factor"], origin: Optional[Term] = None):
         self.factors = tuple(dict.fromkeys(factors))
+        self.origin = origin
         self._factor_key = tuple(factor.expr for factor in sorted(self.factors))
         self._hash = hash(":".join(self._factor_key))
 
@@ -45,7 +51,9 @@ def __eq__(self, other: Any) -> bool:
         if isinstance(other, Term):
             return self._factor_key == other._factor_key
         if isinstance(other, str):
-            return self._factor_key == tuple(sorted(other.split(":")))
+            return self._factor_key == tuple(
+                sorted([m.group("factor") for m in self.FACTOR_MATCHER.finditer(other)])
+            )
         return NotImplemented
 
     def __lt__(self, other: Any) -> bool:
@@ -58,4 +66,4 @@ def __lt__(self, other: Any) -> bool:
         return NotImplemented
 
     def __repr__(self) -> str:
-        return ":".join(factor.expr for factor in self.factors)
+        return ":".join(repr(factor) for factor in self.factors)
diff --git a/formulaic/transforms/lag.py b/formulaic/transforms/lag.py
index 5aae7cbe..257c9320 100644
--- a/formulaic/transforms/lag.py
+++ b/formulaic/transforms/lag.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import functools
-from typing import Any, TYPE_CHECKING
+from typing import Any
 
 import numpy
 import pandas
diff --git a/formulaic/transforms/patsy_compat.py b/formulaic/transforms/patsy_compat.py
index 0831d77e..e5bb5191 100644
--- a/formulaic/transforms/patsy_compat.py
+++ b/formulaic/transforms/patsy_compat.py
@@ -3,12 +3,12 @@
 from formulaic.utils.stateful_transforms import stateful_transform
 
 from .contrasts import (
+    UNSET,
     DiffContrasts,
     HelmertContrasts,
     PolyContrasts,
     SumContrasts,
     TreatmentContrasts,
-    UNSET,
 )
 from .scale import scale
 
diff --git a/pyproject.toml b/pyproject.toml
index 0e2e08d8..697b6c6c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -112,6 +112,7 @@ dependencies = [
 
 [tool.hatch.envs.lint.scripts]
 check = [
+    "ruff check",
     "ruff format --check",
     "mypy formulaic",
 ]


Test patch:
diff --git a/tests/parser/test_parser.py b/tests/parser/test_parser.py
index 02a614dc..d574d0d2 100644
--- a/tests/parser/test_parser.py
+++ b/tests/parser/test_parser.py
@@ -63,6 +63,49 @@
         "lhs": (["a"], ["b"]),
         "rhs": (["1", "c"], ["1", "d"]),
     },
+    # Formula stages
+    "[ a ~ b ]": {"root": ["1", "a_hat"], "deps": ({"lhs": ["a"], "rhs": ["1", "b"]},)},
+    "y ~ [ a ~ b ]": {
+        "lhs": ["y"],
+        "rhs": {"root": ["1", "a_hat"], "deps": ({"lhs": ["a"], "rhs": ["1", "b"]},)},
+    },
+    "[ a ~ [ b ~ c ] ]": {
+        "root": ["1", "a_hat"],
+        "deps": (
+            {
+                "lhs": ["a"],
+                "rhs": {
+                    "root": ["1", "b_hat"],
+                    "deps": ({"lhs": ["b"], "rhs": ["1", "c"]},),
+                },
+            },
+        ),
+    },
+    "y ~ [ a ~ [ b ~ c ] + [d ~ e] ]": {
+        "lhs": ["y"],
+        "rhs": {
+            "root": ["1", "a_hat"],
+            "deps": (
+                {
+                    "lhs": ["a"],
+                    "rhs": {
+                        "root": ["1", "b_hat", "d_hat"],
+                        "deps": (
+                            {"lhs": ["b"], "rhs": ["1", "c"]},
+                            {"lhs": ["d"], "rhs": ["1", "e"]},
+                        ),
+                    },
+                },
+            ),
+        },
+    },
+    "y ~ [ a:b ~ c + d ]": {
+        "lhs": ["y"],
+        "rhs": {
+            "root": ["1", "`a:b_hat`"],
+            "deps": ({"lhs": ["a:b"], "rhs": ["1", "c", "d"]},),
+        },
+    },
     # Products
     "a:b": ["1", "a:b"],
     "b:a + a:b": ["1", "b:a"],
@@ -92,8 +135,10 @@
     "{a | b | c}": ["1", "a | b | c"],
 }
 
-PARSER = DefaultFormulaParser()
-PARSER_NO_INTERCEPT = DefaultFormulaParser(include_intercept=False)
+PARSER = DefaultFormulaParser(feature_flags={"all"})
+PARSER_NO_INTERCEPT = DefaultFormulaParser(
+    include_intercept=False, feature_flags={"all"}
+)
 
 
 class TestFormulaParser:
@@ -226,6 +271,15 @@ def test_feature_flags(self):
         ):
             DefaultFormulaParser().set_feature_flags({}).get_terms("x | y")
 
+    def test_invalid_multistage_formula(self):
+        with pytest.raises(
+            NotImplementedError,
+            match=re.escape(
+                "Nested multistage formulas do not support structured lhs."
+            ),
+        ):
+            DefaultFormulaParser(feature_flags={"all"}).get_terms("[[a ~ b] ~ c]")
+
 
 class TestDefaultOperatorResolver:
     @pytest.fixture
diff --git a/tests/parser/types/test_factor.py b/tests/parser/types/test_factor.py
index b29c1a5e..1140db26 100644
--- a/tests/parser/types/test_factor.py
+++ b/tests/parser/types/test_factor.py
@@ -52,3 +52,4 @@ def test_to_terms(self):
 
     def test_repr(self):
         assert repr(Factor("a")) == "a"
+        assert repr(Factor("a:b")) == "`a:b`"
diff --git a/tests/transforms/test_contrasts.py b/tests/transforms/test_contrasts.py
index 53104bcc..e992ddf4 100644
--- a/tests/transforms/test_contrasts.py
+++ b/tests/transforms/test_contrasts.py
@@ -10,14 +10,14 @@
 from formulaic.materializers import FactorValues
 from formulaic.model_spec import ModelSpec
 from formulaic.transforms.contrasts import (
-    ContrastsRegistry as contr,
+    UNSET,
     ContrastsState,
-)
-from formulaic.transforms.contrasts import (
     SumContrasts,
-    encode_contrasts,
-    UNSET,
     _UnsetSentinel,
+    encode_contrasts,
+)
+from formulaic.transforms.contrasts import (
+    ContrastsRegistry as contr,
 )
 from formulaic.utils.sparse import categorical_encode_series_to_sparse_csc_matrix