Skip to content

Commit 4aba6d0

Browse files
willGraham01willGraham01
authored
Add notebook illustrating how to work with the Continuous Treatment model we've provided (#187)
* Write example out as a docs page so working is explicit * Add new page to site navigation * Start on linear example notebook * Finish section on building the model * How to bulid the regression function * Building the loss function * Write up on how to solve the problems * Remove placeholder section from continuous treatments docs * Revert "Remove placeholder section from continuous treatments docs" This reverts commit 51a244b. * Add reference to example in docs * We now have a new notebook * Organise sidebar page order * Fix hanging nav page * MLP attachments to the treatment model (#200) * Tidy up import block * graph.ricardo -> graph.continuous_treatment * Rename test subdirectory accordingly * MLPAlias and ModelParam go into _types * exmaple_model -> continuous_treatment_model * Add reference to docs page in module docstring * Docstring of example model method * Combine all loose treatment model tests * Rectify the treatment model constructor function * Fix the tikz diagram so it matches the bult model * Hide ugly attachment of MLPs inside constructor function for treatment model * Fix discrepencies in new graph docs * Fix docstring from rebase barf * Remove references to phi_x * Removes the TODO label that is causing ruff to fail. Add manual note to corresponding issue #197 for this to be addressed
1 parent ebab373 commit 4aba6d0

17 files changed

Lines changed: 921 additions & 197 deletions

‎docs/diagrams/continuous-treatment-model.svg‎

Lines changed: 74 additions & 77 deletions
Loading

‎docs/diagrams/continuous-treatment-model.tex‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,22 @@
11
\documentclass[crop,tikz,convert={outext=.svg,command=\unexpanded{pdf2svg \infile\space\outfile}}]{standalone}
22

33
\usepackage{tikz}
4+
\usetikzlibrary{backgrounds}
45

56
\begin{document}
6-
\begin{tikzpicture}[line cap=round,line join=round,line width=1pt,x=2cm,y=2cm]
7+
8+
\begin{tikzpicture}[
9+
line cap=round,
10+
line join=round,
11+
line width=1pt,
12+
x=2cm,
13+
y=2cm,
14+
show background rectangle,
15+
background rectangle/.style={
16+
fill=white,
17+
draw=none,
18+
},
19+
]
720

821
\draw[orange] (-0.2,1) arc(0:180:0.3) -- (-0.8,0) arc(180:270:0.3) -- (2,-0.3) arc (-90:90:0.3) -- (0.1,0.3) arc(270:180:0.3) -- cycle;
922
\node[orange,anchor=west] at (2.3,0) {visible};
@@ -21,17 +34,15 @@
2134
\node at (1,1) {$U_X$};
2235
\node at (2,1) {$U_Y$};
2336
\node at (1.5,1.5) {$C$};
24-
\node at (0.5,0.5) {$\Theta_X$};
2537
\draw[->, shorten <=3mm, shorten >=3mm] (0,0) -- (1,0);
2638
\draw[->, shorten <=3mm, shorten >=3mm] (1,0) -- (2,0);
2739
\draw[->, shorten <=3mm, shorten >=3mm] (1,1) -- (2,1);
2840
\draw[->, shorten <=3mm, shorten >=3mm] (1,1) -- (1,0);
2941
\draw[->, shorten <=3mm, shorten >=3mm] (2,1) -- (2,0);
3042
\draw[->, shorten <=3mm, shorten >=3mm] (1.5,1.5) -- (1,1);
3143
\draw[->, shorten <=3mm, shorten >=3mm] (1.5,1.5) -- (2,1);
32-
\draw[->, shorten <=3mm, shorten >=3mm] (0.5,0.5) -- (1,0);
3344
\draw[->, shorten <=3mm, shorten >=9mm] (-0.5,1) -- (0,0);
34-
\draw[->, shorten <=3mm, shorten >=9mm] (-0.5,1) -- (0.5,0.5);
45+
\draw[->, shorten <=3mm, shorten >=9mm] (-0.5,1) -- (1,0);
3546
\draw[->, shorten <=3mm, shorten >=12mm] (-0.5,1) -- (1,1);
3647

3748
\end{tikzpicture}

‎docs/theory/continuous-treatments-linear-example.md‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
# Continuous Treatment Models: Reduction to Linear Problem with Quadratic Constraint
2+
3+
This document assumes (and uses) the [Continuous Treatment model setup](./continuous-treatments.md), it's notation, etc.
4+
This pages sets up a problem that we can analytically solve, which we then use to demonstrate how to interact with the core components of the framework provided by `causalprog`.
5+
The corresponding notebook can be [found here](https://github.com/UCL/causalprog/blob/main/examples/continuous_treatments_linear.ipynb).
6+
7+
To that end, we will set the following:
8+
9+
<!-- prettier-ignore -->
10+
\begin{align*}
11+
\sigma_{czl} = \frac{1}{2}\mathbb{I}_{d_z}, \quad
12+
\mathcal{D}_{eval} = \left\{ (\tilde{x}, \tilde{z}, \tilde{l} ) \right\}, \quad
13+
g(x, z, l) = -x \mathbb{I}_{d_z}, \quad
14+
\hat{r}(x, z, l) = 0, \quad
15+
f_Y(u_y, x, l) = \frac{\theta_Y}{l}(u_y - x)^2,
16+
\end{align*}
17+
18+
where $\mathbb{I}_{d_z} = \frac{1}{\sqrt{d_z}}(1, 1, 1, ...)^{\top}\in\mathbb{R}^{d_Z}$ and $(\tilde{x}, \tilde{z}, \tilde{l} )$ is some chosen evaluation point (of which, only $\tilde{x}$ will turn out to be relevant).
19+
20+
Note that this effectively forces us to pick constant functions for $f_r$ (constant value $\infty$) and $f_m$ (constant value 0), and thus also gives us $\theta_m$ and $\theta_r$ independent problems.
21+
As we will shortly see, the choice for $f_{\pi}$ will also be irrelevant due to the nature of the problem we are getting up, so for argument's sake it can just map to the constant 1-vector (and the problem is independent of $\theta_{\pi}$ too).
22+
23+
Let us also define $\alpha(x) = \frac{3}{4}(1 + 3x^2) > 0$, which is a constant with respect to the parameters of the problem $\theta$.
24+
We can also immediately deduce that
25+
26+
<!-- prettier-ignore -->
27+
\begin{align*}
28+
m_y = \sigma_{czl}^{\top}g(x, z, l) = -\frac{x}{2}, \quad
29+
v_y = \frac{3}{4}.
30+
\end{align*}
31+
32+
We can then make the following simplifications to the forms of the regression function, learn initialiser, and causal responses:
33+
34+
<!-- prettier-ignore -->
35+
\begin{align*}
36+
r(x, z, l)
37+
& = \int f_Y(u_y, x, l) \sum_{c}\pi_{ul}(c)p_N(u_y; m_y, v_y) \ \mathrm{d}u_y
38+
= \int\frac{\theta_Y}{l}(u_y - x)^2 \sum_{c}\pi_{ul}(c)p_N(u_y; -\frac{x}{2}, \frac{3}{4}) \ \mathrm{d} u_y \\
39+
& = \int\frac{\theta_Y}{l}(u_y - x)^2 p_N(u_y; -\frac{x}{2}, \frac{3}{4}) \ \mathrm{d} u_y
40+
= \frac{\theta_Y}{l}\mathbb{E}\left[(U - x)^2 \ | \ U\sim\mathcal{N}\left(-\frac{x}{2}, \frac{3}{4}\right)\right] \\
41+
& = \frac{\theta_Y}{l}\left(\frac{3}{4} + \left(\frac{-x}{2}\right)^2 - 2x\left(-\frac{x}{2}\right) + x^2\right)
42+
= \frac{\theta_Y \alpha(x)}{l}, \\
43+
B(\theta) & = \frac{1}{n_{eval}}\sum_{\mathcal{D}_{eval}}\left(\hat{r}_i - r(x^{(i)}, z^{(i)}, l^{(i)}) \right)^2
44+
= \frac{\theta_Y^2 \alpha(\tilde{x})^2}{\tilde{l}^2}, \\
45+
d(x, l) & = \int f_Y(u_y, x, l)p_N(u_y; 0, 1) \ \mathrm{d}u_y
46+
= \frac{\theta_Y}{l}(1 + x^2).
47+
\end{align*}
48+
49+
This means that we have $\theta^{\star} = \left\{ \theta_Y = 0 \right\}$, since $B(\theta^{\star}) = 0$.
50+
Note that we have used the fact that $\sum_{c}\pi_{ul}(c) = 1$, since we now have that everything else in the integrand is $c$-independent.
51+
52+
Therefore, given $\delta^2 := \epsilon > 0$, our problem
53+
54+
<!-- prettier-ignore -->
55+
\begin{align*}
56+
\mathrm{min} / \mathrm{max}_{\theta} d(x, l; \theta)
57+
& \quad\text{ subject to }\quad
58+
B(\theta) \leq B(\theta^{\star}) + \epsilon,
59+
\end{align*}
60+
61+
reduces to
62+
63+
<!-- prettier-ignore -->
64+
\begin{align*}
65+
\mathrm{min} / \mathrm{max}_{\theta} \frac{\theta_Y}{l}(1 + x^2)
66+
& \quad\text{ subject to }\quad
67+
\frac{\theta_Y^2 \alpha(\tilde{x})^2}{\tilde{l}^2} \leq \delta^2.
68+
\end{align*}
69+
70+
Furthermore, the constraint is now a simple quadratic in $\theta_Y$ which we can solve for, giving
71+
72+
<!-- prettier-ignore -->
73+
\begin{align*}
74+
-\frac{\delta \tilde{l}}{\alpha(\tilde{x})} \leq \theta_Y \leq \frac{\delta \tilde{l}}{\alpha(\tilde{x})},
75+
\end{align*}
76+
77+
which then lets us immediately read off the solution to the maximisation and minimisation problem:
78+
79+
<!-- prettier-ignore -->
80+
\begin{align*}
81+
\text{max: attained at } \theta_Y = \frac{\delta \tilde{l}}{\alpha(\tilde{x})},
82+
& \quad\text{with objective value } \quad \frac{\delta\tilde{l}}{l}\frac{1 + x^2}{\alpha(\tilde{x})}, \\
83+
\text{min: attained at } \theta_Y = -\frac{\delta l}{\alpha(\tilde{x})},
84+
& \quad\text{with objective value } \quad -\frac{\delta\tilde{l}}{l}\frac{1 + x^2}{\alpha(\tilde{x})}. \\
85+
\end{align*}
86+
87+
Ergo, for given choices of $x$, $\tilde{x}$, $\tilde{l}$, and $\delta$, we now have an analytic solution that we can compare against.
88+
89+
Note that if we were to solve this problem by direct Lagrangian methods, we would construct
90+
91+
<!-- prettier-ignore -->
92+
\begin{align*}
93+
\mathcal{L}(\theta, \lambda) &= \frac{\theta_Y}{l}(1+x^2) - \lambda \max(\frac{\theta_Y^2 \alpha(\tilde{x})^2}{\tilde{l}^2}, 0),
94+
\end{align*}
95+
96+
and shortly deduce that
97+
98+
<!-- prettier-ignore -->
99+
\begin{align*}
100+
\nabla\mathcal{L}(\theta, \lambda) = 0,
101+
& \Leftrightarrow
102+
\theta_Y = \pm\frac{\delta\tilde{l}}{\alpha(\tilde{x})},
103+
\quad
104+
\lambda = \pm \frac{\tilde{l}}{2l} \frac{1+x^2}{\delta \alpha(\tilde{x})}, \\
105+
\text{with} \ d(x, l; \theta) &= \pm\frac{\delta\tilde{l}}{l}\frac{1 + x^2}{\alpha(\tilde{x})}.
106+
\end{align*}

‎docs/users/graph.md‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -158,24 +158,22 @@ of the library and are not used in the current demonstration applications.
158158

159159
## Continuous Treatment models
160160

161-
Many of the examples in causalprog use an example graph, representing a
161+
Many of the examples in `causalprog` use an example graph, representing a
162162
continuous treatment model, for a problem proposed by Ricardo Silva:
163163

164164
![Illustration of the continuous treatment model that we discuss.](../diagrams/continuous-treatment-model.svg)
165165

166-
causalprog provides a helper function to quickly generate this graph. This
167-
function must be passed five required keyword arguments that tell the graph
168-
how to compute the nodes `"u_x"`, `"u_y"`, `"phi_x"`, `"x"` and `"y"` from
169-
their parents.
166+
`causalprog` provides a helper function to quickly generate this graph.
167+
This function must be passed keyword arguments that tell the graph how to compute the nodes `"u_x"`, `"u_y"`, `"x"` and `"y"` from their parents.
168+
Python callables defining $f_r$ and $f_m$ can also be provided.
170169

171170
```python
172-
from causalprog.graph.ricardo import example_model
171+
from causalprog.graph.continuous_treatment import continuous_treatment_model
173172

174-
graph = example_model(
173+
graph = continuous_treatment_model(
175174
compute_u_x=lambda values: values["c"] + 1.0,
176175
compute_u_y=lambda values: values["c"] * 2,
177-
compute_phi_x=lambda values: values["l"],
178-
compute_x=lambda values: values["z"] + values["phi_x"] - values["u_x"],
176+
compute_x=lambda values: values["z"] + values["l"] - values["u_x"],
179177
compute_y=lambda values: values["x"] * values["u_y"],
180178
)
181179
```

0 commit comments

Comments
 (0)