deploy: e09ffa4

Author: TomeHirata

_modules/dte_adj.html

@@ -195,14 +195,14 @@ <h1>Source code for dte_adj</h1><div class="highlight"><pre>
         <span class="n">n_bootstrap</span><span class="p">:</span> <span class="nb">int</span><span class="p">,</span>
     <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;Compute expected DTEs.&quot;&quot;&quot;</span>
-        <span class="n">treatment_cdf</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">treatment_cdf_mat</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_cumulative_distribution</span><span class="p">(</span>
+        <span class="n">treatment_cdf</span><span class="p">,</span> <span class="n">treatment_cdf_mat</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_cumulative_distribution</span><span class="p">(</span>
             <span class="n">target_treatment_arm</span><span class="p">,</span>
             <span class="n">locations</span><span class="p">,</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">covariates</span><span class="p">,</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">treatment_arms</span><span class="p">,</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">outcomes</span><span class="p">,</span>
         <span class="p">)</span>
-        <span class="n">control_cdf</span><span class="p">,</span> <span class="n">_</span><span class="p">,</span> <span class="n">control_cdf_mat</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_cumulative_distribution</span><span class="p">(</span>
+        <span class="n">control_cdf</span><span class="p">,</span> <span class="n">control_cdf_mat</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_cumulative_distribution</span><span class="p">(</span>
             <span class="n">control_treatment_arm</span><span class="p">,</span>
             <span class="n">locations</span><span class="p">,</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">covariates</span><span class="p">,</span>
@@ -665,18 +665,16 @@ <h1>Source code for dte_adj</h1><div class="highlight"><pre>
             <span class="n">binominal</span> <span class="o">=</span> <span class="p">(</span><span class="n">outcomes</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span> <span class="o">&lt;=</span> <span class="n">locations</span><span class="p">)</span> <span class="o">*</span> <span class="mi">1</span>  <span class="c1"># (n_records, n_loc)</span>
             <span class="k">for</span> <span class="n">fold</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">folds</span><span class="p">):</span>
                 <span class="n">fold_mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">folds</span> <span class="o">!=</span> <span class="n">fold</span><span class="p">)</span> <span class="o">&amp;</span> <span class="n">treatment_mask</span>
-                <span class="n">covariates_train</span> <span class="o">=</span> <span class="n">covariates</span><span class="p">[</span><span class="n">fold_mask</span><span class="p">]</span>
-                <span class="n">binominal_train</span> <span class="o">=</span> <span class="n">binominal</span><span class="p">[</span><span class="n">fold_mask</span><span class="p">]</span>
-                <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">binominal_train</span><span class="p">))</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
-                    <span class="bp">self</span><span class="o">.</span><span class="n">model</span> <span class="o">=</span> <span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">base_model</span><span class="p">)</span>
-                    <span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">covariates_train</span><span class="p">,</span> <span class="n">binominal_train</span><span class="p">)</span>
                 <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">s_list</span><span class="p">:</span>
                     <span class="n">s_mask</span> <span class="o">=</span> <span class="n">strata</span> <span class="o">==</span> <span class="n">s</span>
                     <span class="n">weight</span> <span class="o">=</span> <span class="p">(</span><span class="n">s_mask</span> <span class="o">&amp;</span> <span class="n">treatment_mask</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span> <span class="o">/</span> <span class="n">s_mask</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
                     <span class="n">superset_mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">folds</span> <span class="o">==</span> <span class="n">fold</span><span class="p">)</span> <span class="o">&amp;</span> <span class="n">s_mask</span>
                     <span class="n">subset_train_mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">folds</span> <span class="o">!=</span> <span class="n">fold</span><span class="p">)</span> <span class="o">&amp;</span> <span class="n">s_mask</span> <span class="o">&amp;</span> <span class="n">treatment_mask</span>
                     <span class="n">covariates_train</span> <span class="o">=</span> <span class="n">covariates</span><span class="p">[</span><span class="n">subset_train_mask</span><span class="p">]</span>
                     <span class="n">binominal_train</span> <span class="o">=</span> <span class="n">binominal</span><span class="p">[</span><span class="n">subset_train_mask</span><span class="p">]</span>
+                    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">binominal_train</span><span class="p">))</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
+                        <span class="bp">self</span><span class="o">.</span><span class="n">model</span> <span class="o">=</span> <span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">base_model</span><span class="p">)</span>
+                        <span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">covariates_train</span><span class="p">,</span> <span class="n">binominal_train</span><span class="p">)</span>
 
                     <span class="n">pred</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_model_prediction</span><span class="p">(</span>
                         <span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="p">,</span> <span class="n">covariates</span><span class="p">[</span><span class="n">superset_mask</span><span class="p">]</span>
@@ -695,8 +693,9 @@ <h1>Source code for dte_adj</h1><div class="highlight"><pre>
                     <span class="n">fold_mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">folds</span> <span class="o">!=</span> <span class="n">fold</span><span class="p">)</span> <span class="o">&amp;</span> <span class="n">treatment_mask</span>
                     <span class="n">covariates_train</span> <span class="o">=</span> <span class="n">covariates</span><span class="p">[</span><span class="n">fold_mask</span><span class="p">]</span>
                     <span class="n">binominal_train</span> <span class="o">=</span> <span class="n">binominal</span><span class="p">[</span><span class="n">fold_mask</span><span class="p">]</span>
-                    <span class="bp">self</span><span class="o">.</span><span class="n">model</span> <span class="o">=</span> <span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">base_model</span><span class="p">)</span>
+                    <span class="c1"># Pool the records across strata and train the model</span>
                     <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">binominal_train</span><span class="p">))</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
+                        <span class="bp">self</span><span class="o">.</span><span class="n">model</span> <span class="o">=</span> <span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">base_model</span><span class="p">)</span>
                         <span class="bp">self</span><span class="o">.</span><span class="n">model</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">covariates_train</span><span class="p">,</span> <span class="n">binominal_train</span><span class="p">)</span>
                     <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">s_list</span><span class="p">:</span>
                         <span class="n">s_mask</span> <span class="o">=</span> <span class="n">strata</span> <span class="o">==</span> <span class="n">s</span>
@@ -705,7 +704,12 @@ <h1>Source code for dte_adj</h1><div class="highlight"><pre>
                         <span class="n">subset_train_mask</span> <span class="o">=</span> <span class="p">(</span><span class="n">folds</span> <span class="o">!=</span> <span class="n">fold</span><span class="p">)</span> <span class="o">&amp;</span> <span class="n">s_mask</span> <span class="o">&amp;</span> <span class="n">treatment_mask</span>
                         <span class="n">covariates_train</span> <span class="o">=</span> <span class="n">covariates</span><span class="p">[</span><span class="n">subset_train_mask</span><span class="p">]</span>
                         <span class="n">binominal_train</span> <span class="o">=</span> <span class="n">binominal</span><span class="p">[</span><span class="n">subset_train_mask</span><span class="p">]</span>
-                        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">binominal_train</span><span class="p">))</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
+                        <span class="c1"># TODO: revisit the logic here</span>
+                        <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">binominal_train</span><span class="p">))</span> <span class="o">&gt;</span> <span class="mi">1</span><span class="p">:</span>
+                            <span class="c1"># self.model = deepcopy(self.base_model)</span>
+                            <span class="c1"># self.model.fit(covariates_train, binominal_train)</span>
+                            <span class="k">pass</span>
+                        <span class="k">else</span><span class="p">:</span>
                             <span class="n">pred</span> <span class="o">=</span> <span class="n">binominal_train</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
                             <span class="n">superset_prediction</span><span class="p">[</span><span class="n">superset_mask</span><span class="p">,</span> <span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="n">pred</span>
                             <span class="n">prediction</span><span class="p">[</span><span class="n">superset_mask</span><span class="p">,</span> <span class="n">i</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span>