Add trim option to marginal rate computation

openfisca · Oct 10, 2015 · e478fd6 · e478fd6
1 parent 2bf9ac3
commit e478fd6
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/openfisca_core/rates.py b/openfisca_core/rates.py
@@ -25,12 +25,19 @@
 
 from __future__ import division
 
+import numpy
+
 
 def average_rate(target = None, varying = None):
     # target: numerator, varying: denominator
     return 1 - target / (varying * (varying != 0) + (varying == 0))
 
 
-def marginal_rate(target = None, varying = None):
+def marginal_rate(target = None, varying = None, trim = None):
     # target: numerator, varying: denominator
-    return 1 - (target[:-1] - target[1:]) / (varying[:-1] - varying[1:])
+    marginal_rate = 1 - (target[:-1] - target[1:]) / (varying[:-1] - varying[1:])
+    if trim is not None:
+        marginal_rate = numpy.where(marginal_rate <= max(trim), marginal_rate, numpy.nan)
+        marginal_rate = numpy.where(marginal_rate >= min(trim), marginal_rate, numpy.nan)
+
+    return marginal_rate