Deploying to gh-pages from @ f3fae34 🚀

PolicyEngine · Dec 19, 2023 · 1639f6f · 1639f6f
1 parent af0e4a4
commit 1639f6f
Show file tree

Hide file tree

Showing 8 changed files with 102 additions and 9 deletions.
diff --git a/_downloads/621e2c7711e8623173875c33222eddb6/changelog.yaml b/_downloads/621e2c7711e8623173875c33222eddb6/changelog.yaml
@@ -435,3 +435,8 @@
     fixed:
     - Random seed fixed for each simulation.
   date: 2023-12-16 15:34:37
+- bump: patch
+  changes:
+    added:
+    - Tools for branching
+  date: 2023-12-19 11:12:24
diff --git a/_modules/policyengine_core/simulations/simulation.html b/_modules/policyengine_core/simulations/simulation.html
@@ -299,14 +299,22 @@ <h1>Source code for policyengine_core.simulations.simulation</h1><div class="hig
     <span class="n">datasets</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="n">Dataset</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;The list of datasets available for this simulation.&quot;&quot;&quot;</span>
 
+    <span class="n">baseline</span><span class="p">:</span> <span class="s2">&quot;Simulation&quot;</span> <span class="o">=</span> <span class="kc">None</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;The baseline simulation, if this simulation is a reform.&quot;&quot;&quot;</span>
+
+    <span class="n">is_over_dataset</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Whether this simulation is built over a dataset.&quot;&quot;&quot;</span>
+
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
         <span class="n">tax_benefit_system</span><span class="p">:</span> <span class="s2">&quot;TaxBenefitSystem&quot;</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="n">populations</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Population</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="n">situation</span><span class="p">:</span> <span class="nb">dict</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="n">dataset</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Type</span><span class="p">[</span><span class="n">Dataset</span><span class="p">]]</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
         <span class="n">reform</span><span class="p">:</span> <span class="n">Reform</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span>
+        <span class="n">trace</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span>
     <span class="p">):</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">is_over_dataset</span> <span class="o">=</span> <span class="n">dataset</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span>
         <span class="n">reform_applied_after</span> <span class="o">=</span> <span class="kc">False</span>
         <span class="k">if</span> <span class="n">tax_benefit_system</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
             <span class="k">if</span> <span class="p">(</span>
@@ -338,8 +346,10 @@ <h1>Source code for policyengine_core.simulations.simulation</h1><div class="hig
 
         <span class="bp">self</span><span class="o">.</span><span class="n">invalidated_caches</span> <span class="o">=</span> <span class="nb">set</span><span class="p">()</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">debug</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">trace</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">tracer</span><span class="p">:</span> <span class="n">SimpleTracer</span> <span class="o">=</span> <span class="n">SimpleTracer</span><span class="p">()</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">trace</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="n">trace</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">tracer</span><span class="p">:</span> <span class="n">SimpleTracer</span> <span class="o">=</span> <span class="p">(</span>
+            <span class="n">SimpleTracer</span><span class="p">()</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">trace</span> <span class="k">else</span> <span class="n">FullTracer</span><span class="p">()</span>
+        <span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">opt_out_cache</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>
         <span class="c1"># controls the spirals detection; check for performance impact if &gt; 1</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">max_spiral_loops</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">10</span>
@@ -407,6 +417,18 @@ <h1>Source code for policyengine_core.simulations.simulation</h1><div class="hig
             <span class="n">hashed_input</span> <span class="o">=</span> <span class="nb">hash</span><span class="p">(</span><span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">original_input</span><span class="p">))</span> <span class="o">%</span> <span class="mi">1000000</span>
             <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="n">hashed_input</span><span class="p">)</span>
 
+        <span class="k">if</span> <span class="n">reform</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">baseline</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">get_branch</span><span class="p">(</span><span class="s2">&quot;baseline&quot;</span><span class="p">)</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">baseline</span><span class="o">.</span><span class="n">trace</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">trace</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">baseline</span><span class="o">.</span><span class="n">tracer</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">tracer</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">baseline</span><span class="o">.</span><span class="n">tax_benefit_system</span> <span class="o">=</span> <span class="p">(</span>
+                <span class="bp">self</span><span class="o">.</span><span class="n">default_tax_benefit_system_instance</span>
+            <span class="p">)</span>
+        <span class="k">else</span><span class="p">:</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">baseline</span> <span class="o">=</span> <span class="kc">None</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">parent_branch</span> <span class="o">=</span> <span class="kc">None</span>
+
 <div class="viewcode-block" id="Simulation.apply_reform"><a class="viewcode-back" href="../../../python_api/simulations.html#policyengine_core.simulations.simulation.Simulation.apply_reform">[docs]</a>    <span class="k">def</span> <span class="nf">apply_reform</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">reform</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">tuple</span><span class="p">,</span> <span class="n">Reform</span><span class="p">]):</span>
         <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">reform</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
             <span class="k">for</span> <span class="n">subreform</span> <span class="ow">in</span> <span class="n">reform</span><span class="p">:</span>
@@ -587,6 +609,8 @@ <h1>Source code for policyengine_core.simulations.simulation</h1><div class="hig
             <span class="n">variable_name</span><span class="p">,</span> <span class="n">period</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">branch_name</span>
         <span class="p">)</span>
 
+        <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="nb">hash</span><span class="p">(</span><span class="n">variable_name</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">period</span><span class="p">))</span> <span class="o">%</span> <span class="mi">1000000</span><span class="p">)</span>
+
         <span class="k">try</span><span class="p">:</span>
             <span class="n">result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_calculate</span><span class="p">(</span><span class="n">variable_name</span><span class="p">,</span> <span class="n">period</span><span class="p">)</span>
             <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">result</span><span class="p">,</span> <span class="n">EnumArray</span><span class="p">)</span> <span class="ow">and</span> <span class="n">decode_enums</span><span class="p">:</span>
@@ -732,6 +756,18 @@ <h1>Source code for policyengine_core.simulations.simulation</h1><div class="hig
         <span class="k">if</span> <span class="n">cached_array</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
             <span class="k">return</span> <span class="n">cached_array</span>
 
+        <span class="k">if</span> <span class="n">variable</span><span class="o">.</span><span class="n">requires_computation_after</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="k">if</span> <span class="n">variable</span><span class="o">.</span><span class="n">requires_computation_after</span> <span class="ow">not</span> <span class="ow">in</span> <span class="p">[</span>
+                <span class="n">node</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s2">&quot;name&quot;</span><span class="p">)</span> <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">tracer</span><span class="o">.</span><span class="n">stack</span>
+            <span class="p">]:</span>
+                <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
+                    <span class="sa">f</span><span class="s2">&quot;Variable </span><span class="si">{</span><span class="n">variable_name</span><span class="si">}</span><span class="s2"> requires </span><span class="si">{</span><span class="n">variable</span><span class="o">.</span><span class="n">requires_computation_after</span><span class="si">}</span><span class="s2"> to be requested first. The full stack is: &quot;</span>
+                    <span class="o">+</span> <span class="s2">&quot;</span><span class="se">\n</span><span class="s2">&quot;</span><span class="o">.</span><span class="n">join</span><span class="p">(</span>
+                        <span class="sa">f</span><span class="s2">&quot;  - </span><span class="si">{</span><span class="n">node</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;name&#39;</span><span class="p">)</span><span class="si">}</span><span class="s2"> </span><span class="si">{</span><span class="n">node</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;period&#39;</span><span class="p">)</span><span class="si">}</span><span class="s2">, </span><span class="si">{</span><span class="n">node</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">&#39;branch_name&#39;</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span>
+                        <span class="k">for</span> <span class="n">node</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">tracer</span><span class="o">.</span><span class="n">stack</span>
+                    <span class="p">)</span>
+                <span class="p">)</span>
+
         <span class="k">if</span> <span class="n">variable</span><span class="o">.</span><span class="n">definition_period</span> <span class="o">==</span> <span class="n">MONTH</span> <span class="ow">and</span> <span class="n">period</span><span class="o">.</span><span class="n">unit</span> <span class="o">==</span> <span class="n">YEAR</span><span class="p">:</span>
             <span class="k">if</span> <span class="n">variable</span><span class="o">.</span><span class="n">quantity_type</span> <span class="o">==</span> <span class="n">QuantityType</span><span class="o">.</span><span class="n">STOCK</span><span class="p">:</span>
                 <span class="n">contained_months</span> <span class="o">=</span> <span class="n">period</span><span class="o">.</span><span class="n">get_subperiods</span><span class="p">(</span><span class="n">MONTH</span><span class="p">)</span>
@@ -1349,6 +1385,7 @@ <h1>Source code for policyengine_core.simulations.simulation</h1><div class="hig
         <span class="n">branch</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">clone</span><span class="p">(</span><span class="n">clone_tax_benefit_system</span><span class="o">=</span><span class="n">clone_system</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">branches</span><span class="p">[</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">branch</span>
         <span class="n">branch</span><span class="o">.</span><span class="n">branch_name</span> <span class="o">=</span> <span class="n">name</span>
+        <span class="n">branch</span><span class="o">.</span><span class="n">parent_branch</span> <span class="o">=</span> <span class="bp">self</span>
         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">trace</span><span class="p">:</span>
             <span class="n">branch</span><span class="o">.</span><span class="n">trace</span> <span class="o">=</span> <span class="kc">True</span>
             <span class="n">branch</span><span class="o">.</span><span class="n">tracer</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">tracer</span>

diff --git a/_modules/policyengine_core/variables/variable.html b/_modules/policyengine_core/variables/variable.html
@@ -365,6 +365,9 @@ <h1>Source code for policyengine_core.variables.variable</h1><div class="highlig
     <span class="n">hidden_input</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span>
 <span class="w">    </span><span class="sd">&quot;&quot;&quot;Whether the variable is hidden from the input screen entirely on PolicyEngine.&quot;&quot;&quot;</span>
 
+    <span class="n">requires_computation_after</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span>
+<span class="w">    </span><span class="sd">&quot;&quot;&quot;Name of a variable that must be computed before this variable.&quot;&quot;&quot;</span>
+
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">baseline_variable</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">name</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="o">.</span><span class="vm">__name__</span>
         <span class="n">attr</span> <span class="o">=</span> <span class="p">{</span>
@@ -501,6 +504,10 @@ <h1>Source code for policyengine_core.variables.variable</h1><div class="highlig
             <span class="n">attr</span><span class="p">,</span> <span class="s2">&quot;hidden_input&quot;</span><span class="p">,</span> <span class="n">allowed_type</span><span class="o">=</span><span class="nb">bool</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="kc">False</span>
         <span class="p">)</span>
 
+        <span class="bp">self</span><span class="o">.</span><span class="n">requires_computation_after</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">set</span><span class="p">(</span>
+            <span class="n">attr</span><span class="p">,</span> <span class="s2">&quot;requires_computation_after&quot;</span><span class="p">,</span> <span class="n">allowed_type</span><span class="o">=</span><span class="nb">str</span>
+        <span class="p">)</span>
+
         <span class="n">formulas_attr</span><span class="p">,</span> <span class="n">unexpected_attrs</span> <span class="o">=</span> <span class="n">helpers</span><span class="o">.</span><span class="n">_partition</span><span class="p">(</span>
             <span class="n">attr</span><span class="p">,</span>
             <span class="k">lambda</span> <span class="n">name</span><span class="p">,</span> <span class="n">value</span><span class="p">:</span> <span class="n">name</span><span class="o">.</span><span class="n">startswith</span><span class="p">(</span><span class="n">config</span><span class="o">.</span><span class="n">FORMULA_NAME_PREFIX</span><span class="p">),</span>