[ce27e21] | 1 | # TODO: include dispersion docs with the disperser models |
---|
| 2 | from math import sqrt |
---|
[14de349] | 3 | import numpy as np |
---|
[ce27e21] | 4 | from scipy.special import gammaln |
---|
[14de349] | 5 | |
---|
[ce27e21] | 6 | class Dispersion(object): |
---|
| 7 | """ |
---|
| 8 | Base dispersion object. |
---|
| 9 | |
---|
| 10 | Subclasses should define *_weights(center, sigma, lb, ub)* |
---|
| 11 | which returns the x points and their corresponding weights. |
---|
| 12 | """ |
---|
| 13 | type = "base disperser" |
---|
| 14 | default = dict(npts=35, width=0, nsigmas=3) |
---|
| 15 | def __init__(self, npts=None, width=None, nsigmas=None): |
---|
| 16 | self.npts = self.default['npts'] if npts is None else npts |
---|
| 17 | self.width = self.default['width'] if width is None else width |
---|
| 18 | self.nsigmas = self.default['nsigmas'] if nsigmas is None else nsigmas |
---|
[14de349] | 19 | |
---|
| 20 | def get_pars(self): |
---|
[ce27e21] | 21 | pars = {'type': self.type} |
---|
| 22 | pars.update(self.__dict__) |
---|
| 23 | return pars |
---|
| 24 | |
---|
| 25 | def set_weights(self, values, weights): |
---|
| 26 | raise RuntimeError("set_weights is only available for ArrayDispersion") |
---|
| 27 | |
---|
| 28 | def get_weights(self, center, lb, ub, relative): |
---|
| 29 | """ |
---|
| 30 | Return the weights for the distribution. |
---|
| 31 | |
---|
| 32 | *center* is the center of the distribution |
---|
[14de349] | 33 | |
---|
[ce27e21] | 34 | *lb*,*ub* are the min and max allowed values |
---|
| 35 | |
---|
| 36 | *relative* is True if the distribution width is proportional to the |
---|
| 37 | center value instead of absolute. For polydispersity use relative. |
---|
| 38 | For orientation parameters use absolute. |
---|
| 39 | """ |
---|
[14de349] | 40 | npts, width, nsigmas = self.npts, self.width, self.nsigmas |
---|
[ce27e21] | 41 | sigma = self.width * center if relative else self.width |
---|
[14de349] | 42 | if sigma == 0: |
---|
[ce27e21] | 43 | return np.array([center], 'd'), np.array([1.], 'd') |
---|
| 44 | return self._weights(center, sigma, lb, ub) |
---|
| 45 | |
---|
| 46 | def _weights(self, center, sigma, lb, ub): |
---|
| 47 | """actual work of computing the weights""" |
---|
| 48 | raise NotImplementedError |
---|
| 49 | |
---|
| 50 | def _linspace(self, center, sigma, lb, ub): |
---|
| 51 | """helper function to provide linear spaced weight points within range""" |
---|
| 52 | npts, nsigmas = self.npts, self.nsigmas |
---|
| 53 | x = center + np.linspace(-nsigmas*sigma, +nsigmas*sigma, npts) |
---|
| 54 | x = x[(x >= lb) & (x <= ub)] |
---|
| 55 | return x |
---|
| 56 | |
---|
| 57 | |
---|
| 58 | class GaussianDispersion(Dispersion): |
---|
| 59 | type = "gaussian" |
---|
| 60 | default = dict(npts=35, width=0, nsigmas=3) |
---|
| 61 | def _weights(self, center, sigma, lb, ub): |
---|
| 62 | x = self._linspace(center, sigma, lb, ub) |
---|
[14de349] | 63 | px = np.exp((x-center)**2 / (-2.0 * sigma * sigma)) |
---|
| 64 | return x, px |
---|
| 65 | |
---|
[ce27e21] | 66 | |
---|
| 67 | class RectangleDispersion(Dispersion): |
---|
| 68 | type = "rectangle" |
---|
| 69 | default = dict(npts=35, width=0, nsigmas=1.70325) |
---|
| 70 | def _weights(self, center, sigma, lb, ub): |
---|
| 71 | x = self._linspace(center, sigma*sqrt(3.0), lb, ub) |
---|
| 72 | px = np.ones_like(x) |
---|
| 73 | return x, px |
---|
| 74 | |
---|
| 75 | |
---|
| 76 | class LogNormalDispersion(Dispersion): |
---|
| 77 | type = "lognormal" |
---|
| 78 | default = dict(npts=80, width=0, nsigmas=8) |
---|
| 79 | def _weights(self, center, sigma, lb, ub): |
---|
| 80 | x = self._linspace(center, sigma, max(lb,1e-8), max(ub,1e-8)) |
---|
| 81 | px = np.exp(-0.5*(np.log(x)-center)**2)/sigma**2/(x*sigma) |
---|
| 82 | return x, px |
---|
| 83 | |
---|
| 84 | |
---|
| 85 | class SchulzDispersion(Dispersion): |
---|
| 86 | type = "schulz" |
---|
| 87 | default = dict(npts=80, width=0, nsigmas=8) |
---|
| 88 | def _weights(self, center, sigma, lb, ub): |
---|
| 89 | x = self._linspace(center, sigma, max(lb,1e-8), max(ub,1e-8)) |
---|
| 90 | R= x/center |
---|
| 91 | z = (center/sigma)**2 |
---|
| 92 | arg = z*np.log(z) + (z-1)*np.log(R) - R*z - np.log(center) - gammaln(z) |
---|
| 93 | px = np.exp(arg) |
---|
| 94 | return x, px |
---|
| 95 | |
---|
| 96 | |
---|
| 97 | class ArrayDispersion(Dispersion): |
---|
| 98 | type = "array" |
---|
| 99 | default = dict(npts=35, width=0, nsigmas=1) |
---|
| 100 | def __init__(self, npts=None, width=None, nsigmas=None): |
---|
| 101 | Dispersion.__init__(self, npts, width, nsigmas) |
---|
| 102 | self.values = np.array([0.], 'd') |
---|
| 103 | self.weights = np.array([1.], 'd') |
---|
| 104 | |
---|
| 105 | def set_weights(self, values, weights): |
---|
| 106 | self.values = np.ascontiguousarray(values, 'd') |
---|
| 107 | self.weights = np.ascontiguousarray(weights, 'd') |
---|
| 108 | self.npts = len(values) |
---|
| 109 | |
---|
| 110 | def _weights(self, center, sigma, lb, ub): |
---|
| 111 | # TODO: interpolate into the array dispersion using npts |
---|
| 112 | x = center + self.values*sigma |
---|
| 113 | idx = (x>=lb)&(x<=ub) |
---|
| 114 | x = x[idx] |
---|
| 115 | px = self.weights[idx] |
---|
| 116 | return x, px |
---|
| 117 | |
---|
| 118 | |
---|
| 119 | models = dict((d.type,d) for d in ( |
---|
| 120 | GaussianDispersion, RectangleDispersion, |
---|
| 121 | ArrayDispersion, SchulzDispersion, LogNormalDispersion |
---|
| 122 | )) |
---|
| 123 | |
---|
| 124 | def get_weights(disperser, n, width, nsigma, value, limits, relative): |
---|
| 125 | cls = models[disperser] |
---|
| 126 | obj = cls(n, width, nsigma) |
---|
| 127 | v,w = obj.get_weights(value, limits[0], limits[1], relative) |
---|
| 128 | return v,w |
---|