main functions

src/qp/projectors/__init__.py

@@ -1,3 +1,4 @@
 from .projector_base import ProjectorBase
 from .projector_shifts import ProjectorShifts
+from .projector_shifts_widths import ProjectorShiftsWidths
 from .projector_moments import ProjectorMoments

src/qp/projectors/projector_base.py

@@ -1,6 +1,5 @@
 from ..ensemble import Ensemble
 import numpy as np
-from multipledispatch import dispatch
 
 
 class ProjectorBase(object):
@@ -15,45 +14,25 @@ class ProjectorBase(object):
     - get_prior: return the prior distribution of the model given
     the meadured photometric distributions.
     """
-    @dispatch()
-    def __init__(self):
-        self._project_base()
-        self._project()
-
-    @dispatch(np.ndarray, np.ndarray)
-    def __init__(self, zs, pzs):
-        self._project_base(zs, pzs)
-
-    @dispatch(Ensemble)
     def __init__(self, ens):
         self._project_base(ens)
 
-    @dispatch()
-    def _project_base(self):
-        raise NotImplementedError
-
-    @dispatch(np.ndarray, np.ndarray)
-    def _project_base(self, zs, pzs):
-        self.pzs = self._normalize(pzs)
-        self.z = zs
-        self.pz_mean = np.mean(self.pzs, axis=0)
-        self.prior = None
-
-    @dispatch(Ensemble)
     def _project_base(self, ens, z=None):
         if z is None:
             z = np.linspace(0, 1.5, 45)
         self.z = z
-        pzs = ens.pdf(z)
-        pzs = ens.objdata()['pdfs']
-        self.pzs = self._normalize(pzs)
-        self.pz_mean = np.mean(self.pzs, axis=0)
+        nzs = ens.pdf(z)
+        nzs = ens.objdata()['pdfs']
+        self.ens = ens
+        self.nzs = self._normalize(nzs)
+        self.nz_mean = np.mean(self.nzs, axis=0)
+        self.nz_cov = np.cov(self.nzs, rowvar=False)
         self.prior = None
 
-    def _normalize(self, pzs):
-        norms = np.sum(pzs, axis=1)
-        pzs /= norms[:, None]  
-        return pzs
+    def _normalize(self, nzs):
+        norms = np.sum(nzs, axis=1)
+        nzs /= norms[:, None]  
+        return nzs
 
     def evaluate_model(self, *args):
         """
@@ -77,8 +56,16 @@ def sample_prior(self):
         prior = self.get_prior()
         return prior.rvs()
 
-    def save_prior(self):
+    def save_prior(self, mode="dist"):
         """
         Saves the prior distribution to a file.
         """
-        raise NotImplementedError
+        prior = self.get_prior()
+        if mode == "dist":
+            return prior
+        if mode == "file":
+            prior = self.get_prior()
+            np.save("prior_mean.npy", prior.mean)
+            np.save("prior_cov.npy", prior.cov)
+        else:
+            raise NotImplementedError

src/qp/projectors/projector_moments.py

@@ -27,8 +27,8 @@ def __init__(self):
         self._project()
 
     @dispatch(np.ndarray, np.ndarray)
-    def __init__(self, zs, pzs):
-        self._project_base(zs, pzs)
+    def __init__(self, zs, nzs):
+        self._project_base(zs, nzs)
         self._project()
 
     @dispatch(Ensemble)
@@ -37,11 +37,10 @@ def __init__(self, ens):
         self._project()
 
     def _project(self):
-        self.pz_cov = self._get_cov()
-        self.pz_chol = cholesky(self.pz_cov)
+        self.nz_chol = self._get_chol()
 
-    def _get_cov(self):
-        cov = np.cov(self.pzs, rowvar=False)
+    def _get_chol(self):
+        cov = self.nz_cov
         if not self._is_pos_def(cov):
             print('Warning: Covariance matrix is not positive definite')
             print('The covariance matrix will be regularized')
@@ -55,22 +54,23 @@ def _get_cov(self):
                 print('Warning: regularization failed')
                 print('The covariance matrix will be diagonalized')
                 cov = np.diag(np.diag(cov))
-        return cov
+        chol = cholesky(cov)
+        return chol
 
     def _is_pos_def(self, A):
         return np.all(np.linalg.eigvals(A) > 0)
 
-    def evaluate_model(self, pz, alpha):
+    def evaluate_model(self, nz, alpha):
         """
         Samples a photometric distribution 
         from a Gaussian distribution with mean
         and covariance measured from the data.
         """
-        z = pz[0]
-        pz = pz[1]
-        return [z, pz + self.pz_chol @ alpha]
+        z = nz[0]
+        nz = nz[1]
+        return [z, nz + self.nz_chol @ alpha]
 
     def _get_prior(self):
-        return mvn(np.zeros_like(self.pz_mean),
-                   np.ones_like(self.pz_mean))
+        return mvn(np.zeros_like(self.nz_mean),
+                   np.ones_like(self.nz_mean))
 

src/qp/projectors/projector_shifts.py

@@ -1,50 +1,50 @@
 import numpy as np
 from ..ensemble import Ensemble
-from multipledispatch import dispatch
 from .projector_base import ProjectorBase
 from scipy.interpolate import interp1d
-from scipy.stats import multivariate_normal
+from scipy.stats import multivariate_normal as mvn
 
 
 class ProjectorShifts(ProjectorBase):
-    @dispatch()
-    def __init__(self):
-        self._project_base()
-        self._project()
-
-    @dispatch(np.ndarray, np.ndarray)
-    def __init__(self, zs, pzs):
-        self._project_base(zs, pzs)
-        self._project()
-
-    @dispatch(Ensemble)
+    """
+    Projector for the shifts model.
+    The shift model assumes that all the variation in the measured
+    photometric distributions can be described by a single shift in
+    the position of the mean of a fiducial n(z) distribution.
+
+    This shift is calibrated by computing the standard deviations
+    of the measured photometric distributions over redshift.
+    The shift prior is then given by a Gaussian distribution with
+    mean 0 and variance equal to the ratio of the standard deviation
+    of the standard deviations to the mean of the standard deviations.
+    """
     def __init__(self, ens):
         self._project_base(ens)
         self._project()
 
     def _project(self):
         self.shift = self._find_shift()
 
-    def evaluate_model(self, pz, shift):
+    def evaluate_model(self, nz, shift):
         """
         Aplies a shift to the given p(z) distribution.
         This is done by interpolating the p(z) distribution
         at the shifted z values and then evaluating it at the
         original z values.
         """
-        z = pz[0]
-        pz = pz[1]
+        z = nz[0]
+        nz = nz[1]
         z_shift = z + shift
-        pz_shift = interp1d(z_shift, pz,
+        pz_shift = interp1d(z_shift, nz,
                             kind='linear',
                             fill_value='extrapolate')(z)
         return [z, pz_shift]
 
     def _find_shift(self):
-        stds = np.std(self.pzs, axis=1)  # std of each pz
+        stds = np.std(self.nzs, axis=1)  # std of each pz
         s_stds = np.std(stds)            # std of the z-std
         m_stds = np.mean(stds)           # mean of the z-std
         return s_stds / m_stds
 
     def _get_prior(self):
-        return multivariate_normal([0], [self.shift**2])
+        return mvn([0], [self.shift**2])