logprior: add truncated normal

samsam/logprior.py

@@ -18,7 +18,7 @@
 # along with samsam.  If not, see <http://www.gnu.org/licenses/>.
 
 import numpy as np
-from scipy.special import beta as sp_beta
+from scipy.special import beta as sp_beta, erf as sp_erf
 
 log2pi = np.log(2.0 * np.pi)
 halflog2pi = log2pi / 2.0
@@ -44,7 +44,7 @@ def uniform(x, a=0.0, b=1.0):
   """
 
   if x < a or x >= b:
-    raise Exception('uniform: inf')
+    raise Exception('uniform: out of bounds.')
   return (-np.log(b - a))
 
 
@@ -88,7 +88,7 @@ def loguniform(x, loga=0.0, logb=1.0):
   """
 
   if x <= 0:
-    raise Exception('loguniform: inf')
+    raise Exception('loguniform: out of bounds.')
   logx = np.log(x)
   return (-logx + uniform(logx, loga, logb))
 
@@ -117,7 +117,7 @@ def moduniform(x, y, a=0.0, b=1.0):
 
   r = np.sqrt(x**2 + y**2)
   if r <= a or r >= b:
-    raise Exception('moduniform: inf')
+    raise Exception('moduniform: out of bounds.')
   return (-(log2pi + np.log(r * (b - a))))
 
 
@@ -143,6 +143,49 @@ def normal(x, mu=0.0, sig=1.0):
   return (-(halflog2pi + np.log(sig) + ((x - mu) / sig)**2 / 2.0))
 
 
+logtruncnorm_dic = {}
+
+
+def _lazy_logtruncnorm(mu, sig, a, b):
+  r"""
+  Lazy computation of the normalizing coefficient for truncnormal.
+  """
+  if (mu, sig, a, b) not in logtruncnorm_dic:
+    xa = (a - mu) / (np.sqrt(2) * sig)
+    xb = (b - mu) / (np.sqrt(2) * sig)
+    logtruncnorm_dic[(mu, sig, a,
+      b)] = -np.log(np.sqrt(np.pi / 2) * sig * (sp_erf(xb) - sp_erf(xa)))
+  return (logtruncnorm_dic[(mu, sig, a, b)])
+
+
+def truncnormal(x, mu=0.0, sig=1.0, a=0, b=np.inf):
+  r"""
+  Truncated normal distribution with parameters (mu,sig,a,b).
+
+  Parameters
+  ----------
+  x : float
+    Value of the variable.
+  mu : float
+    Mean.
+  sig : float
+    Standard-deviation.
+  a : float
+    Lower bound.
+  b : float
+    Upper bound.
+
+  Returns
+  -------
+  lp : float
+    Log-probability of x.
+  """
+
+  if x < a or x >= b:
+    raise Exception('truncnormal: out of bounds.')
+  return (_lazy_logtruncnorm(mu, sig, a, b) - ((x - mu) / sig)**2 / 2.0)
+
+
 def lognormal(x, mu=0.0, sig=1.0):
   r"""
   Log-normal distribution with parameters (mu,sig).
@@ -163,7 +206,7 @@ def lognormal(x, mu=0.0, sig=1.0):
   """
 
   if x <= 0:
-    raise Exception('lognormal: inf')
+    raise Exception('lognormal: out of bounds.')
   logx = np.log(x)
   return (-logx + normal(logx, mu, sig))
 
@@ -198,7 +241,7 @@ def beta(x, a, b):
   """
 
   if x <= 0 or x >= 1:
-    raise Exception('beta: inf')
+    raise Exception('beta: out of bounds.')
   return ((a - 1.0) * np.log(x) + (b - 1.0) * np.log(1.0 - x) -
     _lazy_logbeta(a, b))
 
@@ -225,6 +268,6 @@ def modbeta(x, y, a, b):
 
   r = np.sqrt(x**2 + y**2)
   if r == 0 or r >= 1:
-    raise Exception('modbeta: inf')
+    raise Exception('modbeta: out of bounds.')
   return ((a - 2.0) * np.log(r) + (b - 1.0) * np.log(1.0 - r) -
     _lazy_logbeta(a, b) - log2pi)