Reorganize and update docs

Author: cthoyt

src/y0/algorithm/identify/id_std.py

@@ -14,6 +14,17 @@ def identify(identification: Identification) -> Expression:
     :param identification: The identification tuple
     :returns: the expression corresponding to the identification
     :raises Unidentifiable: If no appropriate identification can be found
+
+    If you have an instance of a :class:`y0.graph.NxMixedGraph` and a
+    query as an instance of a :class:`y0.dsl.Probability`, use the following:
+
+    .. code-block:: python
+
+        from y0.algorithm.identify import identify, Identification
+        graph = ...
+        query = ...
+        identification = Identification.from_expression(graph=graph, query=query)
+        estimand = identify(identification)
     """
     graph = identification.graph
     treatments = identification.treatments

src/y0/controls.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+
+"""Predicates for good, bad, and neutral controls."""
+
+from .dsl import Probability, Variable
+from .graph import NxMixedGraph
+
+__all__ = [
+    "is_good_control",
+    "is_bad_control",
+]
+
+
+def _control_precondition(graph: NxMixedGraph, query: Probability, variable: Variable):
+    if missing := query.get_variables().difference(graph.nodes()):
+        raise ValueError(f"Query variables missing: {missing}")
+    if variable not in graph.nodes():
+        raise ValueError(f"Test variable missing: {variable}")
+    # TODO does this need to be extended to check that the
+    #  query and variable aren't counterfactual?
+
+
+def is_good_control(graph: NxMixedGraph, query: Probability, variable: Variable) -> bool:
+    """Return if the variable is a good control.
+
+    :param graph: An ADMG
+    :param query: A query in the form of ``P(Y @ X)``
+    :param variable: The variable to check
+    :return: If the variable is a good control
+    """
+    _control_precondition(graph, query, variable)
+
+    raise NotImplementedError
+
+
+def is_bad_control(graph: NxMixedGraph, query: Probability, variable: Variable) -> bool:
+    """Return if the variable is a bad control.
+
+    A bad control is a variable that does not appear in the estimand produced
+    by :func:`y0.algorithm.identify.identify` when applied to a given graph
+    and query.
+
+    :param graph: An ADMG
+    :param query: A query in the form of ``P(Y @ X)``
+    :param variable: The variable to check
+    :return: If the variable is a bad control
+    """
+    _control_precondition(graph, query, variable)
+
+    raise NotImplementedError

src/y0/predicates.py

@@ -2,13 +2,10 @@
 
 """Predicates for expressions."""
 
-from .dsl import Expression, Fraction, Probability, Product, Sum, Variable
-from .graph import NxMixedGraph
+from .dsl import Expression, Fraction, Probability, Product, Sum
 
 __all__ = [
     "has_markov_postcondition",
-    "is_good_control",
-    "is_bad_control",
 ]
 
 
@@ -31,26 +28,3 @@ def has_markov_postcondition(expression: Expression) -> bool:
         )
     else:
         raise TypeError
-
-
-def _control_precondition(graph: NxMixedGraph, query: Probability, variable: Variable):
-    if missing := query.get_variables().difference(graph.nodes()):
-        raise ValueError(f"Query variables missing: {missing}")
-    if variable not in graph.nodes():
-        raise ValueError(f"Test variable missing: {variable}")
-    # TODO does this need to be extended to check that the
-    #  query and variable aren't counterfactual?
-
-
-def is_good_control(graph: NxMixedGraph, query: Probability, variable: Variable) -> bool:
-    """Return if the variable is a good control."""
-    _control_precondition(graph, query, variable)
-
-    raise NotImplementedError
-
-
-def is_bad_control(graph: NxMixedGraph, query: Probability, variable: Variable) -> bool:
-    """Return if the variable is a good control."""
-    _control_precondition(graph, query, variable)
-
-    raise NotImplementedError

tests/test_controls.py

@@ -4,9 +4,9 @@
 
 import unittest
 
+from y0.controls import is_bad_control, is_good_control
 from y0.dsl import U1, U2, M, P, Variable, X, Y, Z
 from y0.graph import NxMixedGraph
-from y0.predicates import is_bad_control, is_good_control
 
 U = Variable("U")
 query = P(Y @ X)