basic truncated normal dist

Author: phobson

paramnormal/dist.py

@@ -904,6 +904,39 @@ def fit(cls, data, **guesses):
         return cls.param_template(R=b*sigma, loc=loc, sigma=sigma)
 
 
+class truncated_normal(BaseDist_Mixin):
+    dist = stats.truncnorm
+    param_template = namedtuple('params', ['lower', 'upper', 'mu', 'sigma'])
+    name = 'truncated normal'
+
+    @staticmethod
+    @utils.greco_deco
+    def _process_args(lower=None, upper=None, mu=None, sigma=None, fit=False):
+        a = None
+        b = None
+        if lower is not None and mu is not None and sigma is not None:
+            a = (lower - mu) / sigma
+
+        if upper is not None and mu is not None and sigma is not None:
+            b = (upper - mu) / sigma
+
+        loc_key, scale_key = utils._get_loc_scale_keys(fit=fit)
+        if fit:
+            akey = 'f0'
+            bkey = 'f1'
+        else:
+            akey = 'a'
+            bkey = 'b'
+        return {akey: a, bkey: b, loc_key: mu, scale_key: sigma}
+
+    @classmethod
+    def fit(cls, data, **guesses):
+        a, b, mu, sigma = cls._fit(data, **guesses)
+        lower = a * sigma + mu
+        upper = b * sigma + mu
+        return cls.param_template(lower=lower, upper=upper, mu=mu, sigma=sigma)
+
+
 __all__ = [
     'normal',
     'lognormal',
@@ -915,4 +948,5 @@ def fit(cls, data, **guesses):
     'pareto',
     'exponential',
     'rice',
+    'truncated_normal',
 ]