1"""
2Addition operator.
3
4Example usage
5-------------
6
7Distribution and a constant::
8
9    >>> distribution = chaospy.Normal(0, 1)+10
10    >>> distribution
11    Add(Normal(mu=0, sigma=1), 10)
12    >>> distribution.sample(5).round(4)
13    array([10.395 ,  8.7997, 11.6476,  9.9553, 11.1382])
14    >>> distribution.fwd([9, 10, 11]).round(4)
15    array([0.1587, 0.5   , 0.8413])
16    >>> distribution.inv(distribution.fwd([9, 10, 11])).round(4)
17    array([ 9., 10., 11.])
18    >>> distribution.pdf([9, 10, 11]).round(4)
19    array([0.242 , 0.3989, 0.242 ])
20    >>> distribution.mom([1, 2, 3]).round(4)
21    array([  10.,  101., 1030.])
22    >>> distribution.ttr([1, 2, 3]).round(4)
23    array([[10., 10., 10.],
24           [ 1.,  2.,  3.]])
25
26Construct joint addition distribution::
27
28    >>> lhs = chaospy.Uniform(2, 3)
29    >>> rhs = chaospy.Uniform(3, 4)
30    >>> addition = lhs + rhs
31    >>> addition
32    Add(Uniform(lower=2, upper=3), Uniform(lower=3, upper=4))
33    >>> joint1 = chaospy.J(lhs, addition)
34    >>> joint2 = chaospy.J(rhs, addition)
35
36Generate random samples::
37
38    >>> joint1.sample(4).round(4)
39    array([[2.2123, 2.0407, 2.3972, 2.2331],
40           [6.0541, 5.2478, 6.1397, 5.6253]])
41    >>> joint2.sample(4).round(4)
42    array([[3.1823, 3.7435, 3.0696, 3.8853],
43           [6.1349, 6.6747, 5.485 , 5.9143]])
44
45Forward transformations::
46
47    >>> lcorr = numpy.array([2.1, 2.5, 2.9])
48    >>> rcorr = numpy.array([3.01, 3.5, 3.99])
49    >>> joint1.fwd([lcorr, lcorr+rcorr]).round(4)
50    array([[0.1 , 0.5 , 0.9 ],
51           [0.01, 0.5 , 0.99]])
52    >>> joint2.fwd([rcorr, lcorr+rcorr]).round(4)
53    array([[0.01, 0.5 , 0.99],
54           [0.1 , 0.5 , 0.9 ]])
55
56Inverse transformations::
57
58    >>> joint1.inv(joint1.fwd([lcorr, lcorr+rcorr])).round(4)
59    array([[2.1 , 2.5 , 2.9 ],
60           [5.11, 6.  , 6.89]])
61    >>> joint2.inv(joint2.fwd([rcorr, lcorr+rcorr])).round(4)
62    array([[3.01, 3.5 , 3.99],
63           [5.11, 6.  , 6.89]])
64
65"""
66from __future__ import division
67from scipy.special import comb
68import numpy
69import chaospy
70
71from ..baseclass import Distribution, OperatorDistribution
72
73
74class Add(OperatorDistribution):
75    """Addition operator."""
76
77    _operator = lambda self, left, right: (left.T+right.T).T
78
79    def __init__(self, left, right):
80        super(Add, self).__init__(
81            left=left,
82            right=right,
83            repr_args=[left, right],
84        )
85
86    def _lower(self, idx, left, right, cache):
87        """
88        Distribution bounds.
89
90        Example:
91            >>> chaospy.Uniform().lower
92            array([0.])
93            >>> chaospy.Add(chaospy.Uniform(), 2).lower
94            array([2.])
95            >>> chaospy.Add(2, chaospy.Uniform()).lower
96            array([2.])
97        """
98        if isinstance(left, Distribution):
99            left = left._get_lower(idx, cache=cache)
100        if isinstance(right, Distribution):
101            right = right._get_lower(idx, cache=cache)
102        return self._operator(left, right)
103
104    def _upper(self, idx, left, right, cache):
105        """
106        Distribution bounds.
107
108        Example:
109            >>> chaospy.Uniform().upper
110            array([1.])
111            >>> chaospy.Add(chaospy.Uniform(), 2).upper
112            array([3.])
113            >>> chaospy.Add(2, chaospy.Uniform()).upper
114            array([3.])
115
116        """
117        if isinstance(left, Distribution):
118            left = left._get_upper(idx, cache=cache)
119        if isinstance(right, Distribution):
120            right = right._get_upper(idx, cache=cache)
121        return self._operator(left, right)
122
123    def _cdf(self, xloc, idx, left, right, cache):
124        if isinstance(right, Distribution):
125            left, right = right, left
126        xloc = (xloc.T-numpy.asfarray(right).T).T
127        uloc = left._get_fwd(xloc, idx, cache=cache)
128        return uloc
129
130    def _pdf(self, xloc, idx, left, right, cache):
131        """
132        Probability density function.
133
134        Example:
135            >>> chaospy.Uniform().pdf([-2, 0, 2, 4])
136            array([0., 1., 0., 0.])
137            >>> chaospy.Add(chaospy.Uniform(), 2).pdf([-2, 0, 2, 4])
138            array([0., 0., 1., 0.])
139            >>> chaospy.Add(2, chaospy.Uniform()).pdf([-2, 0, 2, 4])
140            array([0., 0., 1., 0.])
141
142        """
143        if isinstance(right, Distribution):
144            left, right = right, left
145        xloc = (xloc.T-numpy.asfarray(right).T).T
146        return left._get_pdf(xloc, idx, cache=cache)
147
148    def _ppf(self, uloc, idx, left, right, cache):
149        """
150        Point percentile function.
151
152        Example:
153            >>> chaospy.Uniform().inv([0.1, 0.2, 0.9])
154            array([0.1, 0.2, 0.9])
155            >>> chaospy.Add(chaospy.Uniform(), 2).inv([0.1, 0.2, 0.9])
156            array([2.1, 2.2, 2.9])
157            >>> chaospy.Add(2, chaospy.Uniform()).inv([0.1, 0.2, 0.9])
158            array([2.1, 2.2, 2.9])
159
160        """
161        if isinstance(right, Distribution):
162            left, right = right, left
163        xloc = left._get_inv(uloc, idx, cache=cache)
164        right = numpy.asfarray(right)
165        return self._operator(xloc, right)
166
167    def _mom(self, keys, left, right, cache):
168        """
169        Statistical moments.
170
171        Example:
172            >>> chaospy.Uniform().mom([0, 1, 2, 3]).round(4)
173            array([1.    , 0.5   , 0.3333, 0.25  ])
174            >>> chaospy.Add(chaospy.Uniform(), 2).mom([0, 1, 2, 3]).round(4)
175            array([ 1.    ,  2.5   ,  6.3333, 16.25  ])
176            >>> chaospy.Add(2, chaospy.Uniform()).mom([0, 1, 2, 3]).round(4)
177            array([ 1.    ,  2.5   ,  6.3333, 16.25  ])
178
179        """
180        del cache
181        keys_ = numpy.mgrid[tuple(slice(0, key+1, 1) for key in keys)]
182        keys_ = keys_.reshape(len(self), -1)
183
184        if isinstance(left, Distribution):
185            if chaospy.shares_dependencies(left, right):
186                raise chaospy.StochasticallyDependentError(
187                    "%s: left and right side of sum stochastically dependent." % self)
188            left = [left._get_mom(key) for key in keys_.T]
189        else:
190            left = list(reversed(numpy.array(left).T**keys_.T))
191
192        if isinstance(right, Distribution):
193            right = [right._get_mom(key) for key in keys_.T]
194        else:
195            right = list(reversed(numpy.prod(numpy.array(right).T**keys_.T, -1)))
196
197        out = 0.
198        for idx in range(keys_.shape[1]):
199            key = keys_.T[idx]
200            coef = numpy.prod(comb(keys, key))
201            out += coef*left[idx]*right[idx]*numpy.all(key <= keys)
202        return out
203
204    def _ttr(self, kloc, idx, left, right, cache):
205        """
206        Three terms recurrence coefficients.
207
208        Example:
209            >>> chaospy.Uniform().ttr([0, 1, 2, 3]).round(4)
210            array([[ 0.5   ,  0.5   ,  0.5   ,  0.5   ],
211                   [-0.    ,  0.0833,  0.0667,  0.0643]])
212            >>> chaospy.Add(chaospy.Uniform(), 2).ttr([0, 1, 2, 3]).round(4)
213            array([[ 2.5   ,  2.5   ,  2.5   ,  2.5   ],
214                   [-0.    ,  0.0833,  0.0667,  0.0643]])
215            >>> chaospy.Add(2, chaospy.Uniform()).ttr([0, 1, 2, 3]).round(4)
216            array([[ 2.5   ,  2.5   ,  2.5   ,  2.5   ],
217                   [-0.    ,  0.0833,  0.0667,  0.0643]])
218
219        """
220        del cache
221        if isinstance(right, Distribution):
222            left, right = right, left
223        coeff0, coeff1 = left._get_ttr(kloc, idx)
224        return coeff0+numpy.asarray(right), coeff1
225
226
227def add(left, right):
228    return Add(left, right)
229