deploy: 29de8a0

Author: TomeHirata

_modules/dte_adj/local.html

@@ -30,13 +30,15 @@
           <div class="body" role="main">
 
   <h1>Source code for dte_adj.local</h1><div class="highlight"><pre>
-<span></span><span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
+<span></span><span class="kn">from</span><span class="w"> </span><span class="nn">__future__</span><span class="w"> </span><span class="kn">import</span> <span class="n">annotations</span>
+
+<span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">Tuple</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">dte_adj.stratified</span><span class="w"> </span><span class="kn">import</span> <span class="p">(</span>
     <span class="n">SimpleStratifiedDistributionEstimator</span><span class="p">,</span>
     <span class="n">AdjustedStratifiedDistributionEstimator</span><span class="p">,</span>
 <span class="p">)</span>
-<span class="kn">from</span><span class="w"> </span><span class="nn">dte_adj.util</span><span class="w"> </span><span class="kn">import</span> <span class="n">compute_ldte</span><span class="p">,</span> <span class="n">compute_lpte</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">dte_adj.util</span><span class="w"> </span><span class="kn">import</span> <span class="n">ArrayLike</span><span class="p">,</span> <span class="n">compute_ldte</span><span class="p">,</span> <span class="n">compute_lpte</span><span class="p">,</span> <span class="n">_convert_to_ndarray</span>
 
 
 <div class="viewcode-block" id="SimpleLocalDistributionEstimator">
@@ -64,25 +66,26 @@ <h1>Source code for dte_adj.local</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../api/local.html#dte_adj.SimpleLocalDistributionEstimator.fit">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">fit</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
-        <span class="n">covariates</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">treatment_arms</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">treatment_indicator</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">outcomes</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">strata</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="p">)</span> <span class="o">-&gt;</span> <span class="s2">&quot;SimpleLocalDistributionEstimator&quot;</span><span class="p">:</span>
+        <span class="n">covariates</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">treatment_arms</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">treatment_indicator</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">outcomes</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">strata</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">SimpleLocalDistributionEstimator</span><span class="p">:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Train the SimpleLocalDistributionEstimator.</span>
 
 <span class="sd">        Args:</span>
-<span class="sd">            covariates (np.ndarray): Pre-treatment covariates.</span>
-<span class="sd">            treatment_arms (np.ndarray): Treatment assignment variable (Z).</span>
-<span class="sd">            treatment_indicator (np.ndarray): Treatment indicator variable (D).</span>
-<span class="sd">            outcomes (np.ndarray): Scalar-valued observed outcome.</span>
-<span class="sd">            strata (np.ndarray): Stratum indicators.</span>
+<span class="sd">            covariates: Pre-treatment covariates.</span>
+<span class="sd">            treatment_arms: Treatment assignment variable (Z).</span>
+<span class="sd">            treatment_indicator: Treatment indicator variable (D).</span>
+<span class="sd">            outcomes: Scalar-valued observed outcome.</span>
+<span class="sd">            strata: Stratum indicators.</span>
 
 <span class="sd">        Returns:</span>
 <span class="sd">            SimpleLocalDistributionEstimator: The fitted estimator.</span>
 <span class="sd">        &quot;&quot;&quot;</span>
+        <span class="n">treatment_indicator</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">treatment_indicator</span><span class="p">)</span>
         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">covariates</span><span class="p">,</span> <span class="n">treatment_arms</span><span class="p">,</span> <span class="n">outcomes</span><span class="p">,</span> <span class="n">strata</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">treatment_indicator</span> <span class="o">=</span> <span class="n">treatment_indicator</span>
 
@@ -244,25 +247,26 @@ <h1>Source code for dte_adj.local</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../api/local.html#dte_adj.AdjustedLocalDistributionEstimator.fit">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">fit</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
-        <span class="n">covariates</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">treatment_arms</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">treatment_indicator</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">outcomes</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">strata</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="p">)</span> <span class="o">-&gt;</span> <span class="s2">&quot;AdjustedLocalDistributionEstimator&quot;</span><span class="p">:</span>
+        <span class="n">covariates</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">treatment_arms</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">treatment_indicator</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">outcomes</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+        <span class="n">strata</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">AdjustedLocalDistributionEstimator</span><span class="p">:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Train the AdjustedLocalDistributionEstimator.</span>
 
 <span class="sd">        Args:</span>
-<span class="sd">            covariates (np.ndarray): Pre-treatment covariates.</span>
-<span class="sd">            treatment_arms (np.ndarray): Treatment assignment variable (Z).</span>
-<span class="sd">            treatment_indicator (np.ndarray): Treatment indicator variable (D).</span>
-<span class="sd">            outcomes (np.ndarray): Scalar-valued observed outcome.</span>
-<span class="sd">            strata (np.ndarray): Stratum indicators.</span>
+<span class="sd">            covariates: Pre-treatment covariates.</span>
+<span class="sd">            treatment_arms: Treatment assignment variable (Z).</span>
+<span class="sd">            treatment_indicator: Treatment indicator variable (D).</span>
+<span class="sd">            outcomes: Scalar-valued observed outcome.</span>
+<span class="sd">            strata: Stratum indicators.</span>
 
 <span class="sd">        Returns:</span>
 <span class="sd">            AdjustedLocalDistributionEstimator: The fitted estimator.</span>
 <span class="sd">        &quot;&quot;&quot;</span>
+        <span class="n">treatment_indicator</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">treatment_indicator</span><span class="p">)</span>
         <span class="nb">super</span><span class="p">()</span><span class="o">.</span><span class="n">fit</span><span class="p">(</span><span class="n">covariates</span><span class="p">,</span> <span class="n">treatment_arms</span><span class="p">,</span> <span class="n">outcomes</span><span class="p">,</span> <span class="n">strata</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">treatment_indicator</span> <span class="o">=</span> <span class="n">treatment_indicator</span>
 

_modules/dte_adj/simple.html

@@ -30,11 +30,14 @@
           <div class="body" role="main">
 
   <h1>Source code for dte_adj.simple</h1><div class="highlight"><pre>
-<span></span><span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
+<span></span><span class="kn">from</span><span class="w"> </span><span class="nn">__future__</span><span class="w"> </span><span class="kn">import</span> <span class="n">annotations</span>
+
+<span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
 <span class="kn">from</span><span class="w"> </span><span class="nn">dte_adj.stratified</span><span class="w"> </span><span class="kn">import</span> <span class="p">(</span>
     <span class="n">SimpleStratifiedDistributionEstimator</span><span class="p">,</span>
     <span class="n">AdjustedStratifiedDistributionEstimator</span><span class="p">,</span>
 <span class="p">)</span>
+<span class="kn">from</span><span class="w"> </span><span class="nn">dte_adj.util</span><span class="w"> </span><span class="kn">import</span> <span class="n">ArrayLike</span><span class="p">,</span> <span class="n">_convert_to_ndarray</span>
 
 
 <div class="viewcode-block" id="SimpleDistributionEstimator">
@@ -81,19 +84,23 @@ <h1>Source code for dte_adj.simple</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="SimpleDistributionEstimator.fit">
 <a class="viewcode-back" href="../../api/simple.html#dte_adj.SimpleDistributionEstimator.fit">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">fit</span><span class="p">(</span>
-        <span class="bp">self</span><span class="p">,</span> <span class="n">covariates</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">treatment_arms</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">outcomes</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="o">-&gt;</span> <span class="s2">&quot;SimpleDistributionEstimator&quot;</span><span class="p">:</span>
+        <span class="bp">self</span><span class="p">,</span> <span class="n">covariates</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span> <span class="n">treatment_arms</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span> <span class="n">outcomes</span><span class="p">:</span> <span class="n">ArrayLike</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">SimpleDistributionEstimator</span><span class="p">:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Set parameters.</span>
 
 <span class="sd">        Args:</span>
-<span class="sd">            covariates (np.ndarray): Pre-treatment covariates.</span>
-<span class="sd">            treatment_arms (np.ndarray): The index of the treatment arm.</span>
-<span class="sd">            outcomes (np.ndarray): Scalar-valued observed outcome.</span>
+<span class="sd">            covariates: Pre-treatment covariates.</span>
+<span class="sd">            treatment_arms: The index of the treatment arm.</span>
+<span class="sd">            outcomes: Scalar-valued observed outcome.</span>
 
 <span class="sd">        Returns:</span>
 <span class="sd">            SimpleDistributionEstimator: The fitted estimator.</span>
 <span class="sd">        &quot;&quot;&quot;</span>
+        <span class="n">covariates</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">covariates</span><span class="p">)</span>
+        <span class="n">treatment_arms</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">treatment_arms</span><span class="p">)</span>
+        <span class="n">outcomes</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">outcomes</span><span class="p">)</span>
+
         <span class="k">if</span> <span class="n">covariates</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">treatment_arms</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]:</span>
             <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The shape of covariates and treatment_arm should be same&quot;</span><span class="p">)</span>
 
@@ -147,19 +154,23 @@ <h1>Source code for dte_adj.simple</h1><div class="highlight"><pre>
 <div class="viewcode-block" id="AdjustedDistributionEstimator.fit">
 <a class="viewcode-back" href="../../api/simple.html#dte_adj.AdjustedDistributionEstimator.fit">[docs]</a>
     <span class="k">def</span><span class="w"> </span><span class="nf">fit</span><span class="p">(</span>
-        <span class="bp">self</span><span class="p">,</span> <span class="n">covariates</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">treatment_arms</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">outcomes</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="o">-&gt;</span> <span class="s2">&quot;AdjustedDistributionEstimator&quot;</span><span class="p">:</span>
+        <span class="bp">self</span><span class="p">,</span> <span class="n">covariates</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span> <span class="n">treatment_arms</span><span class="p">:</span> <span class="n">ArrayLike</span><span class="p">,</span> <span class="n">outcomes</span><span class="p">:</span> <span class="n">ArrayLike</span>
+    <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">AdjustedDistributionEstimator</span><span class="p">:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Set parameters.</span>
 
 <span class="sd">        Args:</span>
-<span class="sd">            covariates (np.ndarray): Pre-treatment covariates.</span>
-<span class="sd">            treatment_arms (np.ndarray): The index of the treatment arm.</span>
-<span class="sd">            outcomes (np.ndarray): Scalar-valued observed outcome.</span>
+<span class="sd">            covariates: Pre-treatment covariates.</span>
+<span class="sd">            treatment_arms: The index of the treatment arm.</span>
+<span class="sd">            outcomes: Scalar-valued observed outcome.</span>
 
 <span class="sd">        Returns:</span>
 <span class="sd">            AdjustedDistributionEstimator: The fitted estimator.</span>
 <span class="sd">        &quot;&quot;&quot;</span>
+        <span class="n">covariates</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">covariates</span><span class="p">)</span>
+        <span class="n">treatment_arms</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">treatment_arms</span><span class="p">)</span>
+        <span class="n">outcomes</span> <span class="o">=</span> <span class="n">_convert_to_ndarray</span><span class="p">(</span><span class="n">outcomes</span><span class="p">)</span>
+
         <span class="k">if</span> <span class="n">covariates</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">!=</span> <span class="n">treatment_arms</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]:</span>
             <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;The shape of covariates and treatment_arm should be same&quot;</span><span class="p">)</span>