Repository: matthewwardrop/formulaic

Issue:
Migrating patsy's DesignInto.subset to ModelSpec.?
In preparing statsmodels to move, I have run into an issue where the `subset` method of a `DesignInfo` is used to create a model that is a strict subset of the original model. I can't figure out the best and cleanest way to do this. Any help would be appreciated.

Here is some basic example code for patsy
```python

import numpy as np
import pandas as pd
import patsy

n = 25
df = pd.DataFrame(
    {
        "y": np.random.standard_normal(n),
        "x": np.random.standard_normal(n),
        "z": np.random.standard_normal(n),
        "c": pd.Series(np.random.choice(["a", "b", "d"], size=n),dtype="category"),
    }
)

fmla = "y ~ 1 + x + c + z"
lhs, rhs = patsy.dmatrices(fmla, df)

rhs.design_info.subset(["Intercept","x","z"])
```

which outputs

```
DesignInfo(['Intercept', 'x', 'z'],
           factor_infos={EvalFactor('x'): FactorInfo(factor=EvalFactor('x'),
                                    type='numerical',
                                    state=<factor state>,
                                    num_columns=1),
                         EvalFactor('z'): FactorInfo(factor=EvalFactor('z'),
                                    type='numerical',
                                    state=<factor state>,
                                    num_columns=1)},
           term_codings=OrderedDict([(Term([]),
                                      [SubtermInfo(factors=(),
                                                   contrast_matrices={},
                                                   num_columns=1)]),
                                     (Term([EvalFactor('x')]),
                                      [SubtermInfo(factors=(EvalFactor('x'),),
                                                   contrast_matrices={},
                                                   num_columns=1)]),
                                     (Term([EvalFactor('z')]),
                                      [SubtermInfo(factors=(EvalFactor('z'),),
                                                   contrast_matrices={},
                                                   num_columns=1)])]))
```

This is used in a feature of statsmodels that lets users `drop` columns, and is present in virtually all of the API ([it is in the lowest level base class for usign formulas](https://www.statsmodels.org/stable/dev/generated/statsmodels.base.model.Model.from_formula.html#statsmodels.base.model.Model.from_formula)).

Discussion:
@matthewwardrop: Thanks for surfacing this! I'll take a look tonight. There are some semantics about spanning the intercept that I want to make sure I understand here (if any of these terms includes categorical factors, then simply omitting them from the structure is not necessarily equivalent to building a model matrix as if they were never there). In any case, yes... I'll make sure formulaic can do this too.

Solution patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index a2fd5b67..f5d46f14 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -3,7 +3,7 @@ name: Run Tests
 on:
   push:
     branches:
-      - "**"
+      - "main"
   pull_request:
 
 jobs:
@@ -18,26 +18,26 @@ jobs:
       OS: ${{ matrix.os }}
       PYTHON: ${{ matrix.python-version }}
     steps:
-    - uses: actions/checkout@v1
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install hatch hatch-vcs
-    - name: Run tests
-      run: hatch run tests
-    - if: matrix.python-version == '3.7'
-      name: Check works with min requirements
-      run: hatch run test_min:pytest tests
-    - if: matrix.python-version == '3.9'
-      name: Lint
-      run: hatch run lint:check
-    - name: Upload coverage to Codecov
-      uses: codecov/codecov-action@v1.0.10
-      with:
-        file: ./coverage.xml
-        flags: unittests
-        env_vars: OS,PYTHON
+      - uses: actions/checkout@v1
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install hatch hatch-vcs
+      - name: Run tests
+        run: hatch run tests
+      - if: matrix.python-version == '3.7'
+        name: Check works with min requirements
+        run: hatch run test_min:pytest tests
+      - if: matrix.python-version == '3.9'
+        name: Lint
+        run: hatch run lint:check
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v1.0.10
+        with:
+          file: ./coverage.xml
+          flags: unittests
+          env_vars: OS,PYTHON
diff --git a/docsite/docs/guides/model_specs.ipynb b/docsite/docs/guides/model_specs.ipynb
index 69ecb66f..f5f289f1 100644
--- a/docsite/docs/guides/model_specs.ipynb
+++ b/docsite/docs/guides/model_specs.ipynb
@@ -75,15 +75,19 @@
     "sense to consider on `ModelSpec` instances that are attached to a `ModelMatrix`.\n",
     "They are:\n",
     "\n",
-    "* Property attributes:\n",
+    "* Metadata attributes and methods:\n",
     "    * **column_names**: An ordered sequence of names associated with the columns\n",
     "        of the generated model matrix.\n",
-    "    * **column_indicies**: An ordered mapping from column names to the column \n",
+    "    * **column_indices**: An ordered mapping from column names to the column \n",
     "        index in generated model matrices.\n",
+    "    * **get_column_indicies(...)**: A shorthand method for compiling indices \n",
+    "        for multiple columns from `.column_indices`.\n",
     "    * **terms**: A sequence of `Term` instances that were used to generate this\n",
     "        model matrix.\n",
     "    * **term_indices**: An ordered mapping of `Term` instances to the generated\n",
     "        column indices.\n",
+    "    * **get_term_indices(...)**: A shorthand method for selecting term indices\n",
+    "        from `.term_indices` using formulae.\n",
     "    * **term_slices**: An ordered mapping of `Term` instances to a slice that \n",
     "        when used on the columns of the model matrix will subsample the model\n",
     "        matrix down to those corresponding to each term.\n",
@@ -105,21 +109,26 @@
     "    * **variable_terms**: The reverse lookup of `term_variables`.\n",
     "    * **variable_indices**: A mapping from `Variable` instance to the indices\n",
     "        of the columns in the model matrix that variable.\n",
+    "    * **get_variable_indices(...)**: A shorthand method for compiling indices \n",
+    "        for multiple columns from `.variable_indices`.\n",
     "    * **variables_by_source**: A mapping from source name (typically one of \n",
     "        `\"data\"`, `\"context\"`, or `\"transforms\"`) to the variables derived from\n",
     "        that source.\n",
+    "    * **get_slice(...)**: Build a slice instance that can subset a matrix down \n",
+    "        to the columns associated with a `Term` instance, its string\n",
+    "        representation, a column name, or pre-specified ints/slices.\n",
     "* Utility methods:\n",
     "    * **get_model_matrix(...)**: Build a model matrix using this spec. This \n",
     "        allows a new dataset to be generated using exactly the same encoding \n",
     "        process as an earlier dataset.\n",
     "    * **get_linear_constraints(...)**: Build a set of linear constraints for use\n",
     "        during constrained linear regressions.\n",
-    "    * **get_slice(...)**: Build a slice instance that can subset a matrix down \n",
-    "        to the columns associated with a `Term` instance, its string\n",
-    "        representation, a column name, or pre-specified ints/slices.\n",
     "* Transform methods:\n",
     "    * **update(...)**: Create a copy of this `ModelSpec` instance with the \n",
     "        nominated attributes mutated.\n",
+    "    * **subset(...)**: Create a copy of this `ModelSpec` instance with its\n",
+    "        structure subset to correspond to the form strict subset of terms \n",
+    "        indicted by a formula specification.\n",
     "\n",
     "We'll cover some of these attributes and methods in examples below, but you can\n",
     "always refer to `help(ModelSpec)` for more details."
@@ -773,6 +782,13 @@
     "ModelSpec.from_spec(Formula(lhs=\"y\", rhs=\"a + b\"))"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Some operations, such as `ModelSpec.subset(...)` are also accessible in a mapped way (e.g. via `ModelSpecs.subset(...)`). You can find documentation for the complete set of available methods using `help(ModelSpecs)`."
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",
diff --git a/formulaic/materializers/base.py b/formulaic/materializers/base.py
index deec86a3..4dc18ec1 100644
--- a/formulaic/materializers/base.py
+++ b/formulaic/materializers/base.py
@@ -215,10 +215,19 @@ def _build_model_matrix(
         terms = self._cluster_terms(spec.formula.root, cluster_by=spec.cluster_by)
 
         # Step 1: Determine strategy to maintain structural full-rankness of output matrix
-        scoped_terms_for_terms = self._get_scoped_terms(
-            terms,
-            ensure_full_rank=spec.ensure_full_rank,
-        )
+        # (reusing pre-generated structure if it is available)
+        if spec.structure:
+            scoped_terms_for_terms: Generator[
+                Tuple[Term, Iterable[ScopedTerm]], None, None
+            ] = (
+                (s.term, [st.rehydrate(self.factor_cache) for st in s.scoped_terms])
+                for s in spec.structure
+            )
+        else:
+            scoped_terms_for_terms = self._get_scoped_terms(
+                terms,
+                ensure_full_rank=spec.ensure_full_rank,
+            )
 
         # Step 2: Generate the columns which will be collated into the full matrix
         cols = []
@@ -249,9 +258,7 @@ def _build_model_matrix(
             cols.append((term, scoped_terms, scoped_cols))
 
         # Step 3: Populate remaining model spec fields
-        if spec.structure:
-            cols = self._enforce_structure(cols, spec, drop_rows)  # type: ignore
-        else:
+        if not spec.structure:
             spec = spec.update(
                 structure=[
                     EncodedTermStructure(
@@ -494,9 +501,11 @@ def _simplify_scoped_terms(
                         | (  # type: ignore
                             ScopedTerm(
                                 (
-                                    ScopedFactor(factor_new.factor, reduced=False)
-                                    if factor == factor_new
-                                    else factor
+                                    (
+                                        ScopedFactor(factor_new.factor, reduced=False)
+                                        if factor == factor_new
+                                        else factor
+                                    )
                                     for factor in scoped_term.factors
                                 ),
                                 scale=existing_term.scale * scoped_term.scale,
diff --git a/formulaic/materializers/types/scoped_term.py b/formulaic/materializers/types/scoped_term.py
index 03d4a54b..1912b0e5 100644
--- a/formulaic/materializers/types/scoped_term.py
+++ b/formulaic/materializers/types/scoped_term.py
@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from typing import Any, Iterable, Set
+from typing import Any, Iterable, Mapping, Set
 
+from formulaic.materializers.types.evaluated_factor import EvaluatedFactor
 from formulaic.utils.variables import Variable
 
 from .scoped_factor import ScopedFactor
@@ -62,6 +63,29 @@ def copy(self, *, without_values: bool = False) -> ScopedTerm:
             ]
         return ScopedTerm(factors, scale=self.scale)
 
+    def rehydrate(self, factor_values: Mapping[str, EvaluatedFactor]) -> ScopedTerm:
+        """
+        Rehydrate the `ScopedTerm` instance with new `EvaluatedFactor` instances
+        that have been computed for each factor. An exception will be raised if
+        any required factor is missing from the `factor_values` mapping. This is
+        used to allow the reuse of serialized `ScopedTerm` structure that have
+        had their factor values excised using `.copy(without_values=True)`.
+
+        Args:
+            factor_values: A mapping from factor expressions to `EvaluatedFactor`
+                instances. This mapping should contain all factors in the term.
+        """
+        return ScopedTerm(
+            [
+                ScopedFactor(
+                    factor=factor_values[factor.factor.expr],
+                    reduced=factor.reduced,
+                )
+                for factor in self.factors
+            ],
+            scale=self.scale,
+        )
+
     @property
     def variables(self) -> Set[Variable]:
         return Variable.union(
diff --git a/formulaic/model_spec.py b/formulaic/model_spec.py
index adaeb493..4c66deb7 100644
--- a/formulaic/model_spec.py
+++ b/formulaic/model_spec.py
@@ -11,6 +11,7 @@
     Optional,
     Sequence,
     Set,
+    Tuple,
     Union,
     cast,
 )
@@ -142,10 +143,11 @@ def __post_init__(self) -> None:
 
     # Derived features
 
-    @cached_property
-    def column_names(self) -> Sequence[str]:
+    @property
+    def __structure(self) -> List[EncodedTermStructure]:
         """
-        The names associated with the columns of the generated model matrix.
+        A reference to `.structure` if it is populated, or otherwise an
+        exception is raised.
         """
         if self.structure is None:
             raise RuntimeError(
@@ -153,7 +155,14 @@ def column_names(self) -> Sequence[str]:
                 "likely be resolved by using the `ModelSpec` instance attached "
                 "to the model matrix generated when calling `.get_model_matrix()`."
             )
-        return tuple(feature for row in self.structure for feature in row.columns)
+        return self.structure
+
+    @cached_property
+    def column_names(self) -> Sequence[str]:
+        """
+        The names associated with the columns of the generated model matrix.
+        """
+        return tuple(feature for row in self.__structure for feature in row.columns)
 
     @cached_property
     def column_indices(self) -> Dict[str, int]:
@@ -163,6 +172,19 @@ def column_indices(self) -> Dict[str, int]:
         """
         return {name: i for i, name in enumerate(self.column_names)}
 
+    def get_column_indices(self, columns: Union[str, Sequence[str]]) -> List[int]:
+        """
+        Generate a list of column indices corresponding to the nominated column
+        names. This is useful when you want to slice a model matrix by specific
+        columns, and do not want to have to generate the indices yourself.
+
+        Args:
+            columns: The column names to include in the subset.
+        """
+        if isinstance(columns, str):
+            columns = [columns]
+        return [self.column_indices[column] for column in columns]
+
     @property
     def terms(self) -> List[Term]:
         """
@@ -180,20 +202,45 @@ def term_indices(self) -> Dict[Term, List[int]]:
         up elements of this mapping using the string representation of the
         `Term`.
         """
-        if self.structure is None:
-            raise RuntimeError(
-                "`ModelSpec.structure` has not yet been populated. This will "
-                "likely be resolved by using the `ModelSpec` instance attached "
-                "to the model matrix generated when calling `.get_model_matrix()`."
-            )
         slices = {}
         start = 0
-        for row in self.structure:
+        for row in self.__structure:
             end = start + len(row[2])
             slices[row[0]] = list(range(start, end))
             start = end
         return slices
 
+    def get_term_indices(
+        self, terms_spec: FormulaSpec, **formula_kwargs: Any
+    ) -> List[int]:
+        """
+        Generate a list of column indices corresponding to the columns
+        associated with the nominated `term_spec`.
+
+        This is useful when you want to slice a model matrix by specific terms.
+        If you want to generate new matrices for term subsets, consider using
+        `ModelSpec.subset()` instead.
+
+        The nominated `terms_spec` will be interpreted as a formula
+        specification, and the resulting term set must only include terms that
+        are present in this `ModelSpec` instance. A `ValueError` error will be
+        raised if the `terms_spec` is structured or contains terms not
+        represented by this `ModelSpec`.
+
+        The indices will be ordered according to the order of the terms in the
+        `terms_spec`.
+
+        Args:
+            terms_spec: The specification for the terms for which to extract
+                indices.
+            formula_kwargs: Additional keyword arguments to pass to the
+                `Formula.from_spec` constructor to control (e.g.) ordering.
+        """
+        terms: List[Term] = self.__get_restricted_formula(
+            terms_spec, **formula_kwargs
+        ).root
+        return [idx for term in terms for idx in self.term_indices[term]]
+
     @cached_property
     def term_slices(self) -> Dict[Term, slice]:
         """
@@ -230,14 +277,8 @@ def term_variables(self) -> Dict[Term, Set[Variable]]:
         term. `Variable` instances are enriched strings, with the additional
         attributes `.roles` and `.source`.
         """
-        if self.structure is None:
-            raise RuntimeError(
-                "`ModelSpec.structure` has not yet been populated. This will "
-                "likely be resolved by using the `ModelSpec` instance attached "
-                "to the model matrix generated when calling `.get_model_matrix()`."
-            )
         term_variables = {}
-        for row in self.structure:
+        for row in self.__structure:
             term_variables[row[0]] = Variable.union(
                 *(term.variables for term in row[1]),
             )
@@ -269,16 +310,8 @@ def factor_variables(self) -> Dict[Factor, Set[Variable]]:
         A mapping from `Factor` instances to the variables used in the evaluation
         of that factor.
         """
-
-        if self.structure is None:
-            raise RuntimeError(
-                "`ModelSpec.structure` has not yet been populated. This will "
-                "likely be resolved by using the `ModelSpec` instance attached "
-                "to the model matrix generated when calling `.get_model_matrix()`."
-            )
-
         factor_variables: Dict[Factor, List[Variable]] = defaultdict(list)
-        for s in self.structure:
+        for s in self.__structure:
             for scoped_term in s.scoped_terms:
                 for scoped_factor in scoped_term.factors:
                     factor_variables[scoped_factor.factor.factor].extend(
@@ -354,6 +387,24 @@ def variable_indices(self) -> Dict[Variable, List[int]]:
             for variable, terms in self.variable_terms.items()
         }
 
+    def get_variable_indices(
+        self, variables: Sequence[Union[str, Variable]]
+    ) -> List[int]:
+        """
+        Generate a list of column indices corresponding to the columns associated
+        with the nominated variables. This is useful when you want to slice a model
+        matrix by specific variables, and do not want to have to generate the indices
+        yourself.
+
+        Args:
+            variables: The variable names to include in the subset.
+        """
+        return [
+            idx
+            for variable in variables
+            for idx in self.variable_indices[variable]  # type: ignore # Variables are strings too
+        ]
+
     @cached_property
     def variables_by_source(self) -> Dict[Optional[str], Set[Variable]]:
         """
@@ -366,33 +417,44 @@ def variables_by_source(self) -> Dict[Optional[str], Set[Variable]]:
             variables_by_source[variable.source].add(variable)
         return dict(variables_by_source)
 
-    # Transforms
-
-    def update(self, **kwargs: Any) -> ModelSpec:
-        """
-        Create a copy of this `ModelSpec` instance with the nominated attributes
-        mutated.
+    def get_slice(self, columns_identifier: Union[int, str, Term, slice]) -> slice:
         """
-        return replace(self, **kwargs)
+        Generate a `slice` instance corresponding to the columns associated with
+        the nominated `columns_identifier`. While this is provided for
+        convenience, it is usually better in library code to directly use the
+        indexing metadata methods/attributes associated with the nominated
+        identifier.
 
-    def differentiate(self, *wrt: str, use_sympy: bool = False) -> ModelSpec:
+        Args:
+            columns_identifier: The identifier for which the slice should be
+                generated. Can be one of:
+                    - an integer specifying a specific column index.
+                    - a `Term` instance
+                    - a string representation of a term
+                    - a column name
         """
-        EXPERIMENTAL: Take the gradient of this model spec. When used a linear
-        regression, evaluating a trained model on model matrices generated by
-        this formula is equivalent to estimating the gradient of that fitted
-        form with respect to `vars`.
+        if isinstance(columns_identifier, slice):
+            return columns_identifier
+        if isinstance(columns_identifier, int):
+            return slice(columns_identifier, columns_identifier + 1)
 
-        Args:
-            wrt: The variables with respect to which the gradient should be
-                taken.
-            use_sympy: Whether to use sympy to perform symbolic differentiation.
+        term_slices = self.term_slices
+        if isinstance(columns_identifier, Term):
+            if columns_identifier not in term_slices:
+                raise ValueError(
+                    f"Model matrices built using this spec do not include term: `{columns_identifier}`."
+                )
+            return term_slices[columns_identifier]
+        if columns_identifier in term_slices:
+            return term_slices[columns_identifier]  # type: ignore # Terms hash equivalent to their string repr
 
-        Notes:
-            This method is provisional and may be removed in any future major
-            version.
-        """
-        return self.update(
-            formula=self.formula.differentiate(*wrt, use_sympy=use_sympy),
+        column_indices = self.column_indices
+        if columns_identifier in column_indices:
+            idx = column_indices[columns_identifier]
+            return slice(idx, idx + 1)
+
+        raise ValueError(
+            f"Model matrices built using this spec do not have any columns related to: `{repr(columns_identifier)}`."
         )
 
     # Utility methods
@@ -441,41 +503,77 @@ def get_linear_constraints(self, spec: LinearConstraintSpec) -> LinearConstraint
         """
         return LinearConstraints.from_spec(spec, variable_names=self.column_names)
 
-    def get_slice(self, columns_identifier: Union[int, str, Term, slice]) -> slice:
+    # Transforms
+
+    def update(self, **kwargs: Any) -> ModelSpec:
         """
-        Generate a `slice` instance corresponding to the columns associated with
-        the nominated `columns_identifier`.
+        Create a copy of this `ModelSpec` instance with the nominated attributes
+        mutated.
+        """
+        return replace(self, **kwargs)
+
+    def subset(self, terms_spec: FormulaSpec, **formula_kwargs: Any) -> ModelSpec:
+        """
+        Subset this `ModelSpec` instance to only include the columns associated
+        with the nominated terms.
+
+        This is useful when you want to fit restricted models on a strict subset
+        of features included in this `ModelSpec` instance, and want to generate
+        new model matrices with just these terms.  If you just want to subset an
+        existing model matrix, you can use `ModelSpec.get_term_indices()`
+        instead.
+
+        Terms are selected from this `ModelSpec` instance by constructing a
+        `Formula` instance from the provided `terms_spec`, and then matching the
+        terms with those found in this `ModelSpec` instance. An error will be
+        raised if the `terms_spec` is incompatibly structured or contains terms
+        not represented by this `ModelSpec` instance. The model spec column
+        ordering will follow the ordering of the terms in `terms_spec`.
+
+        Note that subsetting this `ModelSpec` is in general not equivalent to
+        constructing this `ModelSpec` instance from scratch with the provided
+        formula specification. Specifically, the output is likely not to be
+        structurally full-rank whenever categorical variables are involved.
+        Instead, columns generated from the subset model spec are guaranteed to
+        match the corresponding columns generated from this parent model spec.
 
         Args:
-            columns_identifier: The identifier for which the slice should be
-                generated. Can be one of:
-                    - an integer specifying a specific column index.
-                    - a `Term` instance
-                    - a string representation of a term
-                    - a column name
+            terms_spec: The terms to include in the subset. A `Formula` instance
+                will be constructed from this specification, and the resulting
+                terms will be used to select the terms to include from this
+                model spec.
+            formula_kwargs: Additional keyword arguments to pass to the
+                `Formula.from_spec` constructor to control (e.g.) ordering.
         """
-        if isinstance(columns_identifier, slice):
-            return columns_identifier
-        if isinstance(columns_identifier, int):
-            return slice(columns_identifier, columns_identifier + 1)
 
-        term_slices = self.term_slices
-        if isinstance(columns_identifier, Term):
-            if columns_identifier not in term_slices:
-                raise ValueError(
-                    f"Model matrices built using this spec do not include term: `{columns_identifier}`."
-                )
-            return term_slices[columns_identifier]
-        if columns_identifier in term_slices:
-            return term_slices[columns_identifier]  # type: ignore # Terms hash equivalent to their string repr
+        formula: Formula = self.__get_restricted_formula(terms_spec, **formula_kwargs)
+        terms: List[Term] = formula.root
+        terms_set: Set[Term] = set(terms)
+        term_structure = {s.term: s for s in self.__structure if s.term in terms_set}
 
-        column_indices = self.column_indices
-        if columns_identifier in column_indices:
-            idx = column_indices[columns_identifier]
-            return slice(idx, idx + 1)
+        return self.update(
+            formula=formula,
+            structure=[term_structure[term] for term in terms],
+        )
 
-        raise ValueError(
-            f"Model matrices built using this spec do not have any columns related to: `{repr(columns_identifier)}`."
+    def differentiate(self, *wrt: str, use_sympy: bool = False) -> ModelSpec:
+        """
+        EXPERIMENTAL: Take the gradient of this model spec. When used a linear
+        regression, evaluating a trained model on model matrices generated by
+        this formula is equivalent to estimating the gradient of that fitted
+        form with respect to `vars`.
+
+        Args:
+            wrt: The variables with respect to which the gradient should be
+                taken.
+            use_sympy: Whether to use sympy to perform symbolic differentiation.
+
+        Notes:
+            This method is provisional and may be removed in any future major
+            version.
+        """
+        return self.update(
+            formula=self.formula.differentiate(*wrt, use_sympy=use_sympy),
         )
 
     # Only include dataclass fields when pickling.
@@ -484,6 +582,39 @@ def __getstate__(self) -> Dict[str, Any]:
             k: v for k, v in self.__dict__.items() if k in self.__dataclass_fields__
         }
 
+    # Helpers
+
+    def __get_restricted_formula(
+        self, spec: FormulaSpec, **formula_kwargs: Any
+    ) -> Formula:
+        """
+        Construct a `Formula` instance from the provided `spec` that is a
+        restriction of the formula associated with this `ModelSpec` instance. A
+        `ValueError` is raised if the provided `spec` results in a structured
+        formula, or if it contains terms that are not described by this
+        `ModelSpec`.
+
+        Args:
+            spec: The formula specification for the restricted formula.
+            formula_kwargs: Additional keyword arguments to pass to the
+                `Formula.from_spec` constructor to control (e.g.) ordering.
+        """
+        formula = Formula.from_spec(spec, **formula_kwargs)
+        if formula._has_structure:
+            raise ValueError(
+                "Cannot subset a `ModelSpec` using a formula that has structure."
+            )
+
+        terms = formula.root
+
+        missing_terms: Set[Term] = set(terms).difference(self.terms)
+        if missing_terms:
+            raise ValueError(
+                f"Cannot subset a model spec with terms not present in the original model spec: {missing_terms}."
+            )
+
+        return formula
+
 
 class ModelSpecs(Structured[ModelSpec]):
     """
@@ -576,6 +707,37 @@ def get_model_matrix(
             ),
         )
 
+    def subset(self, terms_spec: FormulaSpec) -> ModelSpecs:
+        """
+        Subset this `ModelSpecs` instance to only include the columns associated
+        with the nominated terms. The structure of `terms_spec` must match the
+        structure of this `ModelSpecs` instance where they overlap. For more
+        details, please reference to `ModelSpec.subset`.
+
+        Args:
+            terms_spec: The terms to include in the subset. A `Formula` instance
+                will be constructed from this specification, and the resulting
+                terms will be used to select the terms to include from this
+                model spec.
+        """
+
+        formula = Formula.from_spec(terms_spec)
+
+        def map_formula_structure_onto_model_spec(
+            terms: List[Term], context: Tuple[Union[int, str], ...]
+        ) -> ModelSpec:
+            try:
+                return self[context].subset(terms)
+            except KeyError:
+                raise ValueError(
+                    f"Cannot subset a `ModelSpecs` instance using a formula with a different structure [indexing path `{context}` not found]."
+                )
+
+        return cast(
+            ModelSpecs,
+            formula._map(map_formula_structure_onto_model_spec, as_type=ModelSpecs),
+        )
+
     def differentiate(self, *wrt: str, use_sympy: Any = False) -> ModelSpecs:
         """
         This method proxies the experimental `ModelSpec.differentiate(...)` API.
diff --git a/formulaic/parser/types/structured.py b/formulaic/parser/types/structured.py
index 2d1e0f01..3398d1c7 100644
--- a/formulaic/parser/types/structured.py
+++ b/formulaic/parser/types/structured.py
@@ -145,9 +145,13 @@ def _has_structure(self) -> bool:
 
     def _map(
         self,
-        func: Callable[[ItemType], Any],
+        func: Union[
+            Callable[[ItemType], Any],
+            Callable[[ItemType, Tuple[Union[str, int], ...]], Any],
+        ],
         recurse: bool = True,
         as_type: Optional[Type[Structured]] = None,
+        _context: Tuple[Union[str, int], ...] = (),
     ) -> Structured[Any]:
         """
         Map a callable object onto all the structured objects, returning a
@@ -171,15 +175,21 @@ def _map(
             but with all objects transformed under `func`.
         """
 
-        def apply_func(obj: Any) -> Any:
+        def apply_func(obj: Any, context: Tuple[Union[str, int], ...]) -> Any:
             if recurse and isinstance(obj, Structured):
-                return obj._map(func, recurse=True, as_type=as_type)
+                return obj._map(func, recurse=True, as_type=as_type, _context=context)
             if isinstance(obj, tuple):
-                return tuple(apply_func(o) for o in obj)
-            return func(obj)
+                return tuple(apply_func(o, context + (i,)) for i, o in enumerate(obj))
+            try:
+                return func(obj, context)  # type: ignore
+            except TypeError:
+                return func(obj)  # type: ignore
 
         return (as_type or Structured)(
-            **{key: apply_func(obj) for key, obj in self._structure.items()}
+            **{
+                key: apply_func(obj, _context + (key,))
+                for key, obj in self._structure.items()
+            }
         )
 
     def _flatten(self) -> Generator[ItemType, None, None]:
@@ -436,6 +446,15 @@ def __getitem__(self, key: Any) -> Any:
             return self.root
         if isinstance(key, str) and not key.startswith("_") and key in self._structure:
             return self._structure[key]
+        if isinstance(key, tuple) and len(key) >= 1 and key[0] in self._structure:
+            obj = self[key[0]]
+            if len(key) == 1:
+                return obj
+            if isinstance(obj, Structured):
+                return obj[key[1:]]
+            raise KeyError(
+                f"{key} extends beyond structure of `{self.__class__.__name__}`."
+            )
         raise KeyError(
             f"This `{self.__class__.__name__}` instance does not have structure @ `{repr(key)}`."
         )
@@ -448,6 +467,16 @@ def __setitem__(self, key: Any, value: Any) -> Any:
                 "Substructure keys cannot start with an underscore. "
                 f"The invalid keys are: {set(key for key in self._structure if key.startswith('_'))}."
             )
+        if isinstance(key, tuple) and len(key) > 1 and key[0] in self._structure:
+            obj = self[key[0]]
+            if isinstance(obj, Structured):
+                obj[key[1:]] = value
+                return
+            raise KeyError(
+                f"{key} extends beyond structure of `{self.__class__.__name__}`."
+            )
+        if isinstance(key, tuple) and len(key) == 1:
+            key = key[0]
         self._structure[key] = self.__prepare_item(key, value)
 
     def __iter__(self) -> Generator[Any, None, None]:


Test patch:
diff --git a/tests/materializers/types/test_scoped_term.py b/tests/materializers/types/test_scoped_term.py
index 5696609e..8a947e67 100644
--- a/tests/materializers/types/test_scoped_term.py
+++ b/tests/materializers/types/test_scoped_term.py
@@ -42,5 +42,21 @@ def test_copy(self, scoped_term):
         assert copied.factors == scoped_term.factors
         assert copied.scale == scoped_term.scale
 
+        copied_without_values = scoped_term.copy(without_values=True)
+        assert all(
+            factor.factor.values is None for factor in copied_without_values.factors
+        )
+
+    def test_rehydrate(self, scoped_term):
+        factors = {
+            "a": EvaluatedFactor(Factor("a"), values=[1], variables={Variable("a")}),
+            "b": EvaluatedFactor(Factor("b"), values=[1], variables={Variable("b")}),
+        }
+        rehydrated = scoped_term.copy(without_values=True).rehydrate(factors)
+        assert all(
+            factors[factor.factor.expr] is factor.factor
+            for factor in rehydrated.factors
+        )
+
     def test_variables(self, scoped_term):
         assert scoped_term.variables == {"a", "b"}
diff --git a/tests/test_model_spec.py b/tests/test_model_spec.py
index d19a6c09..f97d325e 100644
--- a/tests/test_model_spec.py
+++ b/tests/test_model_spec.py
@@ -107,28 +107,21 @@ def test_attributes(self, model_spec):
         }
         assert model_spec.variables_by_source == {"data": {"a", "A"}}
 
-    def test_get_model_matrix(self, model_spec, data2):
-        m = model_spec.get_model_matrix(data2)
-
-        assert isinstance(m, pandas.DataFrame)
-        assert model_spec.column_names == tuple(m.columns)
-
-        model_spec = model_spec.update(materializer=None)
-        m2 = model_spec.get_model_matrix(data2)
-        assert isinstance(m2, pandas.DataFrame)
-        assert model_spec.column_names == tuple(m2.columns)
-
-        m3 = model_spec.get_model_matrix(data2, output="sparse")
-        assert isinstance(m3, scipy.sparse.spmatrix)
-
-    def test_get_linear_constraints(self, model_spec):
-        lc = model_spec.get_linear_constraints("`A[T.b]` - a = 3")
-        assert numpy.allclose(lc.constraint_matrix, [[0.0, -1.0, 1.0, 0, 0.0, 0.0]])
-        assert lc.constraint_values == [3]
-        assert lc.variable_names == model_spec.column_names
-
-    def test_differentiate(self, model_spec, formula):
-        assert model_spec.differentiate("a").formula == formula.differentiate("a")
+    def test_get_column_indices(self, model_spec):
+        assert model_spec.get_column_indices("a") == [1]
+        assert model_spec.get_column_indices("A[T.b]:a") == [4]
+        assert model_spec.get_column_indices(["a", "A[T.b]:a"]) == [1, 4]
+
+    def test_get_term_indices(self, model_spec):
+        assert model_spec.get_term_indices("a") == [0, 1]
+        assert model_spec.get_term_indices(["a"]) == [1]
+        assert model_spec.get_term_indices("a + a:A") == [0, 1, 4, 5]
+        assert model_spec.get_term_indices(["a", "a:A"]) == [1, 4, 5]
+        assert model_spec.get_term_indices(["a:A", "a"], ordering="none") == [4, 5, 1]
+
+    def test_get_variable_indices(self, model_spec):
+        assert model_spec.get_variable_indices("a") == [1, 4, 5]
+        assert model_spec.get_variable_indices("A") == [2, 3, 4, 5]
 
     def test_get_slice(self, model_spec):
         s = slice(0, 1)
@@ -154,47 +147,28 @@ def test_get_slice(self, model_spec):
         ):
             model_spec.get_slice("missing")
 
-    def test_model_specs(self, model_spec, data2):
-        model_specs = ModelSpecs(a=model_spec)
-
-        assert isinstance(model_specs.get_model_matrix(data2), ModelMatrices)
-        assert isinstance(model_specs.get_model_matrix(data2).a, ModelMatrix)
-        assert numpy.all(
-            model_specs.get_model_matrix(data2).a == model_spec.get_model_matrix(data2)
-        )
-        sparse_matrices = model_specs.get_model_matrix(data2, output="sparse")
-        assert isinstance(sparse_matrices, ModelMatrices)
-        assert isinstance(sparse_matrices.a, scipy.sparse.spmatrix)
-
-        # Validate missing materializer and output type behaviour
-        model_specs2 = ModelSpecs(
-            lhs=ModelSpec(formula="A"), rhs=ModelSpec(formula="a")
-        )
-        assert numpy.all(
-            model_specs2.get_model_matrix(data2).lhs
-            == ModelSpec(formula="A").get_model_matrix(data2)
-        )
+    def test_get_model_matrix(self, model_spec, data2):
+        m = model_spec.get_model_matrix(data2)
 
-        # Validate non-joint generation of model matrices
-        class MyPandasMaterializer(PandasMaterializer):
-            REGISTER_NAME = "my_pandas_materializer"
+        assert isinstance(m, pandas.DataFrame)
+        assert model_spec.column_names == tuple(m.columns)
 
-        incompatible_specs = ModelSpecs(
-            a=model_spec, b=model_spec.update(materializer=MyPandasMaterializer)
-        )
-        matrices = incompatible_specs.get_model_matrix(data2)
-        assert isinstance(matrices, ModelMatrices)
-        assert numpy.all(matrices.a == model_spec.get_model_matrix(data2))
-        assert numpy.all(matrices.b == model_spec.get_model_matrix(data2))
+        model_spec = model_spec.update(materializer=None)
+        m2 = model_spec.get_model_matrix(data2)
+        assert isinstance(m2, pandas.DataFrame)
+        assert model_spec.column_names == tuple(m2.columns)
 
-        FormulaMaterializerMeta.REGISTERED_NAMES.pop("my_pandas_materializer")
+        m3 = model_spec.get_model_matrix(data2, output="sparse")
+        assert isinstance(m3, scipy.sparse.spmatrix)
 
-        # Validate differentiation
-        assert model_specs.differentiate("a").a == model_spec.differentiate("a")
+    def test_get_linear_constraints(self, model_spec):
+        lc = model_spec.get_linear_constraints("`A[T.b]` - a = 3")
+        assert numpy.allclose(lc.constraint_matrix, [[0.0, -1.0, 1.0, 0, 0.0, 0.0]])
+        assert lc.constraint_values == [3]
+        assert lc.variable_names == model_spec.column_names
 
-        # Validate invalid type checking
-        with pytest.raises(TypeError, match="`ModelSpecs` instances expect all.*"):
-            ModelSpecs("invalid type!")
+    def test_differentiate(self, model_spec, formula):
+        assert model_spec.differentiate("a").formula == formula.differentiate("a")
 
     def test_empty(self):
         model_spec = ModelSpec([])
@@ -256,3 +230,96 @@ def test_unrepresented_term(self):
             "1": slice(0, 1),
             "A": slice(0, 0),
         }
+
+    def test_subset(self, model_spec, data2):
+        subset = model_spec.subset(["a"])
+        assert subset.formula.root == ["a"]
+        assert subset.variables == {"a"}
+        assert numpy.allclose(
+            model_spec.get_model_matrix(data2).values[:, model_spec.get_slice("a")],
+            subset.get_model_matrix(data2).values,
+        )
+
+        subset = model_spec.subset(["a:A", "A"])
+        assert subset.formula.root == ["A", "a:A"]
+        assert subset.variables == {"a", "A"}
+        assert numpy.allclose(
+            model_spec.get_model_matrix(data2).values[
+                :, model_spec.get_term_indices(["A", "a:A"])
+            ],
+            subset.get_model_matrix(data2).values,
+        )
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "Cannot subset a `ModelSpec` using a formula that has structure."
+            ),
+        ):
+            model_spec.subset("a | b")
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                r"Cannot subset a model spec with terms not present in the original model spec: {new_var}"
+            ),
+        ):
+            model_spec.subset("new_var")
+
+    def test_model_specs(self, model_spec, data2):
+        model_specs = ModelSpecs(a=model_spec)
+
+        assert isinstance(model_specs.get_model_matrix(data2), ModelMatrices)
+        assert isinstance(model_specs.get_model_matrix(data2).a, ModelMatrix)
+        assert numpy.all(
+            model_specs.get_model_matrix(data2).a == model_spec.get_model_matrix(data2)
+        )
+        sparse_matrices = model_specs.get_model_matrix(data2, output="sparse")
+        assert isinstance(sparse_matrices, ModelMatrices)
+        assert isinstance(sparse_matrices.a, scipy.sparse.spmatrix)
+
+        # Validate missing materializer and output type behaviour
+        model_specs2 = ModelSpecs(
+            lhs=ModelSpec(formula="A"), rhs=ModelSpec(formula="a")
+        )
+        assert numpy.all(
+            model_specs2.get_model_matrix(data2).lhs
+            == ModelSpec(formula="A").get_model_matrix(data2)
+        )
+
+        # Validate non-joint generation of model matrices
+        class MyPandasMaterializer(PandasMaterializer):
+            REGISTER_NAME = "my_pandas_materializer"
+
+        incompatible_specs = ModelSpecs(
+            a=model_spec, b=model_spec.update(materializer=MyPandasMaterializer)
+        )
+        matrices = incompatible_specs.get_model_matrix(data2)
+        assert isinstance(matrices, ModelMatrices)
+        assert numpy.all(matrices.a == model_spec.get_model_matrix(data2))
+        assert numpy.all(matrices.b == model_spec.get_model_matrix(data2))
+
+        FormulaMaterializerMeta.REGISTERED_NAMES.pop("my_pandas_materializer")
+
+        # Validate differentiation
+        assert model_specs.differentiate("a").a == model_spec.differentiate("a")
+
+        # Validate invalid type checking
+        with pytest.raises(TypeError, match="`ModelSpecs` instances expect all.*"):
+            ModelSpecs("invalid type!")
+
+    def test_model_specs_subset(self, data, data2):
+        ms = Formula("a ~ A + A:a").get_model_matrix(data).model_spec
+
+        mss = ms.subset("a ~ A")
+
+        assert mss.lhs.formula.root == ["a"]
+        assert mss.rhs.formula.root == ["1", "A"]
+
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "Cannot subset a `ModelSpecs` instance using a formula with a different structure [indexing path `('rhs', 0)` not found]."
+            ),
+        ):
+            ms.subset("a ~ A | A:a")