merge attempt

ChangeLog

 Changes since NIFTy 5:
 
+MPI parallelisation over samples in MetricGaussianKL
+====================================================
+
+The classes `MetricGaussianKL` and `MetricGaussianKL_MPI` have been unified
+into one `MetricGaussianKL` class which has MPI support built in.
+
+New approach for random number generation
+=========================================
+
+The code now uses `numpy`'s new `SeedSequence` and `Generator` classes for the
+production of random numbers (introduced in numpy 1.17. This greatly simplifies
+the generation of reproducible random numbers in the presence of MPI parallelism
+and leads to cleaner code overall. Please see the documentation of
+`nifty6.random` for details.
+
 Interface Change for non-linear Operators
 =========================================
 
@@ -17,7 +32,6 @@ behaviour since both `Operator._check_input()` and
 `extra.check_jacobian_consistency()` tests for the new conditions to be
 fulfilled.
 
-
 Special functions for complete Field reduction operations
 =========================================================
 
@@ -66,12 +80,3 @@ User-visible changes:
   replaced by a single function called `makeField`
 - the property `local_shape` has been removed from `Domain` (and subclasses)
   and `DomainTuple`.
-
-Transfer of MPI parallelization into operators:
-===============================================
-
-As was already the case with the `MetricGaussianKL_MPI` in NIFTy5, MPI
-parallelization in NIFTy6 is handled by specialized MPI-enabled operators.
-
-They are accessible via the `nifty6.mpi` namespace, from which they can be
-imported directly: `from nifty6.mpi import MPIenabledOperator`.

demos/Wiener_Filter.ipynb

@@ -140,7 +140,6 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "np.random.seed(40)\n",
     "import nifty6 as ift\n",
     "import matplotlib.pyplot as plt\n",
     "%matplotlib inline"

demos/bench_gridder.py

@@ -5,8 +5,6 @@ import numpy as np
 
 import nifty6 as ift
 
-np.random.seed(40)
-
 N0s, a0s, b0s, c0s = [], [], [], []
 
 for ii in range(10, 26):
@@ -15,15 +13,16 @@ for ii in range(10, 26):
     N = int(2**ii)
     print('N = {}'.format(N))
 
-    uv = np.random.rand(N, 2) - 0.5
-    vis = np.random.randn(N) + 1j*np.random.randn(N)
+    rng = ift.random.current_rng()
+    uv = rng.uniform(-.5, .5, (N,2))
+    vis = rng.normal(0., 1., N) + 1j*rng.normal(0., 1., N)
 
     uvspace = ift.RGSpace((nu, nv))
 
     visspace = ift.TMP_UnstructuredSpace(N)
 
-    img = np.random.randn(nu*nv)
-    img = img.reshape((nu, nv))
+    img = rng.standard_normal((nu, nv))
+    img = ift.makeField(uvspace, img)
     img = ift.makeTMP_fld(uvspace, img)
 
     t0 = time()

demos/bernoulli_demo.py

@@ -27,8 +27,6 @@ import numpy as np
 import nifty6 as ift
 
 if __name__ == '__main__':
-    np.random.seed(41)
-
     # Set up the position space of the signal
     mode = 2
     if mode == 0:
@@ -62,7 +60,7 @@ if __name__ == '__main__':
     p = R(sky)
     mock_position = ift.from_random('normal', harmonic_space)
     tmp = p(mock_position).val.astype(np.float64)
-    data = np.random.binomial(1, tmp)
+    data = ift.random.current_rng().binomial(1, tmp)
     data = ift.TMP_fld.from_raw(R.target, data)
 
     # Compute likelihood and Hamiltonian

demos/getting_started_1.py

@@ -46,8 +46,6 @@ def make_random_mask():
 
 
 if __name__ == '__main__':
-    np.random.seed(42)
-
     # Choose space on which the signal field is defined
     if len(sys.argv) == 2:
         mode = int(sys.argv[1])

demos/getting_started_2.py

@@ -44,8 +44,6 @@ def exposure_2d():
 
 
 if __name__ == '__main__':
-    np.random.seed(42)
-
     # Choose space on which the signal field is defined
     if len(sys.argv) == 2:
         mode = int(sys.argv[1])
@@ -94,7 +92,7 @@ if __name__ == '__main__':
     lamb = R(sky)
     mock_position = ift.from_random('normal', domain)
     data = lamb(mock_position)
-    data = np.random.poisson(data.val.astype(np.float64))
+    data = ift.random.current_rng().poisson(data.val.astype(np.float64))
     data = ift.TMP_fld.from_raw(d_space, data)
     likelihood = ift.PoissonianEnergy(data) @ lamb
 

demos/getting_started_3.py

@@ -33,20 +33,18 @@ import nifty6 as ift
 
 
 def random_los(n_los):
-    starts = list(np.random.uniform(0, 1, (n_los, 2)).T)
-    ends = list(np.random.uniform(0, 1, (n_los, 2)).T)
+    starts = list(ift.random.current_rng().random((n_los, 2)).T)
+    ends = list(ift.random.current_rng().random((n_los, 2)).T)
     return starts, ends
 
 
 def radial_los(n_los):
-    starts = list(np.random.uniform(0, 1, (n_los, 2)).T)
-    ends = list(0.5 + 0*np.random.uniform(0, 1, (n_los, 2)).T)
+    starts = list(ift.random.current_rng().random((n_los, 2)).T)
+    ends = list(0.5 + 0*ift.random.current_rng().random((n_los, 2)).T)
     return starts, ends
 
 
 if __name__ == '__main__':
-    np.random.seed(420)
-
     # Choose between random line-of-sight response (mode=0) and radial lines
     # of sight (mode=1)
     if len(sys.argv) == 2:

demos/getting_started_mf.py

@@ -44,20 +44,18 @@ class SingleTMP_Space(ift.LinearOperator):
 
 
 def random_los(n_los):
-    starts = list(np.random.uniform(0, 1, (n_los, 2)).T)
-    ends = list(np.random.uniform(0, 1, (n_los, 2)).T)
+    starts = list(ift.random.current_rng().random((n_los, 2)).T)
+    ends = list(ift.random.current_rng().random((n_los, 2)).T)
     return starts, ends
 
 
 def radial_los(n_los):
-    starts = list(np.random.uniform(0, 1, (n_los, 2)).T)
-    ends = list(0.5 + 0*np.random.uniform(0, 1, (n_los, 2)).T)
+    starts = list(ift.random.current_rng().random((n_los, 2)).T)
+    ends = list(0.5 + 0*ift.random.current_rng().random((n_los, 2)).T)
     return starts, ends
 
 
 if __name__ == '__main__':
-    np.random.seed(43)
-
     # Choose between random line-of-sight response (mode=0) and radial lines
     # of sight (mode=1)
     if len(sys.argv) == 2:

demos/polynomial_fit.py

@@ -19,7 +19,6 @@ import matplotlib.pyplot as plt
 import numpy as np
 
 import nifty6 as ift
-np.random.seed(12)
 
 
 def polynomial(coefficients, sampling_points):
@@ -86,7 +85,7 @@ if __name__ == '__main__':
     N_params = 10
     N_samples = 100
     size = (12,)
-    x = np.random.random(size) * 10
+    x = ift.random.current_rng().random(size) * 10
     y = np.sin(x**2) * x**3
     var = np.full_like(y, y.var() / 10)
     var[-2] *= 4

nifty6/__init__.py

 from .version import __version__
 
+from . import random
+
 from .domains.domain import TMP_Space
 from .domains.structured_domain import TMP_StructuredSpace
 from .domains.unstructured_domain import TMP_UnstructuredSpace

nifty6/extra.py

@@ -278,17 +278,21 @@ def check_jacobian_consistency(op, loc, tol=1e-8, ntries=100, perf_check=True):
         dir = loc2-loc
         locnext = loc2
         dirnorm = dir.norm()
+        hist = []
         for i in range(50):
             locmid = loc + 0.5*dir
             linmid = op(Linearization.make_var(locmid))
             dirder = linmid.jac(dir)
             numgrad = (lin2.val-lin.val)
             xtol = tol * dirder.norm() / np.sqrt(dirder.size)
+            hist.append((numgrad-dirder).norm())
+#            print(len(hist),hist[-1])
             if (abs(numgrad-dirder) <= xtol).s_all():
                 break
             dir = dir*0.5
             dirnorm *= 0.5
             loc2, lin2 = locmid, linmid
         else:
+            print(hist)
             raise ValueError("gradient and value seem inconsistent")
         loc = locnext

nifty6/library/special_distributions.py

@@ -25,6 +25,7 @@ from ..field import TMP_fld
 from ..linearization import Linearization
 from ..operators.operator import Operator
 from ..sugar import makeOp
+from .. import random
 
 
 def _f_on_np(f, arr):
@@ -67,7 +68,8 @@ class _InterpolationOperator(Operator):
         if table_func is not None:
             if inv_table_func is None:
                 raise ValueError
-            a = func(np.random.randn(10))
+# MR FIXME: not sure whether we should have this in production code
+            a = func(random.current_rng().random(10))
             a1 = _f_on_np(lambda x: inv_table_func(table_func(x)), a)
             np.testing.assert_allclose(a, a1)
             self._table = _f_on_np(table_func, self._table)

nifty6/minimization/metric_gaussian_kl.py

@@ -11,19 +11,66 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #
-# Copyright(C) 2013-2019 Max-Planck-Society
+# Copyright(C) 2013-2020 Max-Planck-Society
 #
 # NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
+
 import numpy as np
 
-from .. import utilities
+from .. import random, utilities
+from ..field import TMP_fld
 from ..linearization import Linearization
+from ..multi_field import TMP_Field
+from ..operators.endomorphic_operator import EndomorphicOperator
 from ..operators.energy_operators import StandardHamiltonian
 from ..probing import approximation2endo
-from ..sugar import makeOp
+from ..sugar import full, makeOp
 from .energy import Energy
 
 
+def _shareRange(nwork, nshares, myshare):
+    nbase = nwork//nshares
+    additional = nwork % nshares
+    lo = myshare*nbase + min(myshare, additional)
+    hi = lo + nbase + int(myshare < additional)
+    return lo, hi
+
+
+def _np_allreduce_sum(comm, arr):
+    if comm is None:
+        return arr
+    from mpi4py import MPI
+    arr = np.array(arr)
+    res = np.empty_like(arr)
+    comm.Allreduce(arr, res, MPI.SUM)
+    return res
+
+
+def _allreduce_sum_field(comm, fld):
+    if comm is None:
+        return fld
+    if isinstance(fld, TMP_fld):
+        return TMP_fld(fld.domain, _np_allreduce_sum(fld.val))
+    res = tuple(
+        TMP_fld(f.domain, _np_allreduce_sum(comm, f.val))
+        for f in fld.values())
+    return TMP_Field(fld.domain, res)
+
+
+class _KLMetric(EndomorphicOperator):
+    def __init__(self, KL):
+        self._KL = KL
+        self._capability = self.TIMES | self.ADJOINT_TIMES
+        self._domain = KL.position.domain
+
+    def apply(self, x, mode):
+        self._check_input(x, mode)
+        return self._KL.apply_metric(x)
+
+    def draw_sample(self, from_inverse=False, dtype=np.float64):
+        return self._KL._metric_sample(from_inverse, dtype)
+
+
 class MetricGaussianKL(Energy):
     """Provides the sampled Kullback-Leibler divergence between a distribution
     and a Metric Gaussian.
@@ -38,6 +85,7 @@ class MetricGaussianKL(Energy):
     typically nonlinear structure of the true distribution these samples have
     to be updated eventually by intantiating `MetricGaussianKL` again. For the
     true probability distribution the standard parametrization is assumed.
+    The samples of this class can be distributed among MPI tasks.
 
     Parameters
     ----------
@@ -62,6 +110,17 @@ class MetricGaussianKL(Energy):
     napprox : int
         Number of samples for computing preconditioner for sampling. No
         preconditioning is done by default.
+    comm : MPI communicator or None
+        If not None, samples will be distributed as evenly as possible
+        across this communicator. If `mirror_samples` is set, then a sample and
+        its mirror image will always reside on the same task.
+    lh_sampling_dtype : type
+        Determines which dtype in data space shall be used for drawing samples
+        from the metric. If the inference is based on complex data,
+        lh_sampling_dtype shall be set to complex accordingly. The reason for
+        the presence of this parameter is that metric of the likelihood energy
+        is just an `Operator` which does not know anything about the dtype of
+        the fields on which it acts. Default is float64.
     _samples : None
         Only a parameter for internal uses. Typically not to be set by users.
 
@@ -79,7 +138,8 @@ class MetricGaussianKL(Energy):
 
     def __init__(self, mean, hamiltonian, n_samples, constants=[],
                  point_estimates=[], mirror_samples=False,
-                 napprox=0, _samples=None, lh_sampling_dtype=np.float64):
+                 napprox=0, comm=None, _samples=None,
+                 lh_sampling_dtype=np.float64):
         super(MetricGaussianKL, self).__init__(mean)
 
         if not isinstance(hamiltonian, StandardHamiltonian):
@@ -88,47 +148,72 @@ class MetricGaussianKL(Energy):
             raise ValueError
         if not isinstance(n_samples, int):
             raise TypeError
-        self._constants = list(constants)
-        self._point_estimates = list(point_estimates)
+        self._constants = tuple(constants)
+        self._point_estimates = tuple(point_estimates)
         if not isinstance(mirror_samples, bool):
             raise TypeError
 
         self._hamiltonian = hamiltonian
 
+        self._n_samples = int(n_samples)
+        if comm is not None:
+            self._comm = comm
+            ntask = self._comm.Get_size()
+            rank = self._comm.Get_rank()
+            self._lo, self._hi = _shareRange(self._n_samples, ntask, rank)
+        else:
+            self._comm = None
+            self._lo, self._hi = 0, self._n_samples
+
+        self._mirror_samples = bool(mirror_samples)
+        self._n_eff_samples = self._n_samples
+        if self._mirror_samples:
+            self._n_eff_samples *= 2
+
         if _samples is None:
             met = hamiltonian(Linearization.make_partial_var(
-                mean, point_estimates, True)).metric
-            if napprox > 1:
+                mean, self._point_estimates, True)).metric
+            if napprox >= 1:
                 met._approximation = makeOp(approximation2endo(met, napprox))
-            _samples = tuple(met.draw_sample(from_inverse=True,
-                                             dtype=lh_sampling_dtype)
-                             for _ in range(n_samples))
-            if mirror_samples:
-                _samples += tuple(-s for s in _samples)
+            _samples = []
+            sseq = random.spawn_sseq(self._n_samples)
+            for i in range(self._lo, self._hi):
+                random.push_sseq(sseq[i])
+                _samples.append(met.draw_sample(from_inverse=True,
+                                                dtype=lh_sampling_dtype))
+                random.pop_sseq()
+            _samples = tuple(_samples)
+        else:
+            if len(_samples) != self._hi-self._lo:
+                raise ValueError("# of samples mismatch")
         self._samples = _samples
-
-        # FIXME Use simplify for constant input instead
-        self._lin = Linearization.make_partial_var(mean, constants)
+        self._lin = Linearization.make_partial_var(mean, self._constants)
         v, g = None, None
-        for s in self._samples:
-            tmp = self._hamiltonian(self._lin+s)
-            if v is None:
-                v = tmp.val.val[()]
-                g = tmp.gradient
-            else:
-                v += tmp.val.val[()]
-                g = g + tmp.gradient
-        self._val = v / len(self._samples)
-        self._grad = g * (1./len(self._samples))
+        if len(self._samples) == 0:  # hack if there are too many MPI tasks
+            tmp = self._hamiltonian(self._lin)
+            v = 0. * tmp.val.val
+            g = 0. * tmp.gradient
+        else:
+            for s in self._samples:
+                tmp = self._hamiltonian(self._lin+s)
+                if self._mirror_samples:
+                    tmp = tmp + self._hamiltonian(self._lin-s)
+                if v is None:
+                    v = tmp.val.val_rw()
+                    g = tmp.gradient
+                else:
+                    v += tmp.val.val
+                    g = g + tmp.gradient
+        self._val = _np_allreduce_sum(self._comm, v)[()] / self._n_eff_samples
+        self._grad = _allreduce_sum_field(self._comm, g) / self._n_eff_samples
         self._metric = None
-        self._napprox = napprox
         self._sampdt = lh_sampling_dtype
 
     def at(self, position):
-        return MetricGaussianKL(position, self._hamiltonian, 0,
-                                self._constants, self._point_estimates,
-                                napprox=self._napprox, _samples=self._samples,
-                                lh_sampling_dtype=self._sampdt)
+        return MetricGaussianKL(
+            position, self._hamiltonian, self._n_samples, self._constants,
+            self._point_estimates, self._mirror_samples, comm=self._comm,
+            _samples=self._samples, lh_sampling_dtype=self._sampdt)
 
     @property
     def value(self):
@@ -139,30 +224,48 @@ class MetricGaussianKL(Energy):
         return self._grad
 
     def _get_metric(self):
+        lin = self._lin.with_want_metric()
         if self._metric is None:
-            lin = self._lin.with_want_metric()
-            mymap = map(lambda v: self._hamiltonian(lin+v).metric,
-                        self._samples)
-            self._unscaled_metric = utilities.my_sum(mymap)
-            self._metric = self._unscaled_metric.scale(1./len(self._samples))
-
-    def unscaled_metric(self):
-        self._get_metric()
-        return self._unscaled_metric, 1/len(self._samples)
+            if len(self._samples) == 0:  # hack if there are too many MPI tasks
+                self._metric = self._hamiltonian(lin).metric.scale(0.)
+            else:
+                mymap = map(lambda v: self._hamiltonian(lin+v).metric,
+                            self._samples)
+                self.unscaled_metric = utilities.my_sum(mymap)
+                self._metric = self.unscaled_metric.scale(1./self._n_eff_samples)
 
     def apply_metric(self, x):
         self._get_metric()
-        return self._metric(x)
+        return _allreduce_sum_field(self._comm, self._metric(x))
 
     @property
     def metric(self):
-        self._get_metric()
-        return self._metric
+        return _KLMetric(self)
 
     @property
     def samples(self):
-        return self._samples
+        if self._comm is not None:
+            res = self._comm.allgather(self._samples)
+            res = [item for sublist in res for item in sublist]
+        else:
+            res = self._samples
+        if self._mirror_samples:
+            res = res + tuple(-item for item in res)
+        return res
+
+    def _unscaled_metric_sample(self, from_inverse=False, dtype=np.float64):
+        if from_inverse:
+            raise NotImplementedError()
+        lin = self._lin.with_want_metric()
+        samp = full(self._hamiltonian.domain, 0.)
+        sseq = random.spawn_sseq(self._n_samples)
+        for i, v in enumerate(self._samples):
+            random.push_sseq(sseq[self._lo+i])
+            samp = samp + self._hamiltonian(lin+v).metric.draw_sample(from_inverse=False, dtype=dtype)
+            if self._mirror_samples:
+                samp = samp + self._hamiltonian(lin-v).metric.draw_sample(from_inverse=False, dtype=dtype)
+            random.pop_sseq()
+        return _allreduce_sum_field(self._comm, samp)
 
-    def __repr__(self):
-        return 'KL ({} samples):\n'.format(len(
-            self._samples)) + utilities.indent(self._hamiltonian.__repr__())
+    def _metric_sample(self, from_inverse=False, dtype=np.float64):
+        return self._unscaled_metric_sample(from_inverse, dtype)/self._n_eff_samples

nifty6/minimization/metric_gaussian_kl_mpi.py

-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <http://www.gnu.org/licenses/>.
-#
-# Copyright(C) 2013-2019 Max-Planck-Society
-#
-# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
-
-from .. import utilities
-from ..linearization import Linearization
-from ..operators.energy_operators import StandardHamiltonian
-from ..operators.endomorphic_operator import EndomorphicOperator
-from .energy import Energy
-from mpi4py import MPI
-import numpy as np
-from ..probing import approximation2endo
-from ..sugar import makeOp, full
-from ..field import TMP_fld
-from ..multi_field import TMP_Field
-
-
-_comm = MPI.COMM_WORLD
-ntask = _comm.Get_size()
-rank = _comm.Get_rank()
-master = (rank == 0)
-
-
-def _shareRange(nwork, nshares, myshare):
-    nbase = nwork//nshares
-    additional = nwork % nshares
-    lo = myshare*nbase + min(myshare, additional)
-    hi = lo + nbase + int(myshare < additional)
-    return lo, hi
-
-
-def np_allreduce_sum(arr):
-    arr = np.array(arr)
-    res = np.empty_like(arr)
-    _comm.Allreduce(arr, res, MPI.SUM)
-    return res
-
-
-def allreduce_sum_field(fld):
-    if isinstance(fld, TMP_fld):
-        return TMP_fld(fld.domain, np_allreduce_sum(fld.val))
-    res = tuple(
-        TMP_fld(f.domain, np_allreduce_sum(f.val))
-        for f in fld.values())
-    return TMP_Field(fld.domain, res)
-
-
-class KLMetric(EndomorphicOperator):
-    def __init__(self, KL):
-        self._KL = KL
-        self._capability = self.TIMES | self.ADJOINT_TIMES
-        self._domain = KL.position.domain
-
-    def apply(self, x, mode):
-        self._check_input(x, mode)
-        return self._KL.apply_metric(x)
-
-    def draw_sample(self, from_inverse=False, dtype=np.float64):
-        return self._KL.metric_sample(from_inverse, dtype)
-
-
-class MetricGaussianKL_MPI(Energy):
-    """Provides the sampled Kullback-Leibler divergence between a distribution
-    and a Metric Gaussian.
-
-    A Metric Gaussian is used to approximate another probability distribution.
-    It is a Gaussian distribution that uses the Fisher information metric of
-    the other distribution at the location of its mean to approximate the
-    variance. In order to infer the mean, a stochastic estimate of the
-    Kullback-Leibler divergence is minimized. This estimate is obtained by
-    sampling the Metric Gaussian at the current mean. During minimization
-    these samples are kept constant; only the mean is updated. Due to the
-    typically nonlinear structure of the true distribution these samples have
-    to be updated eventually by intantiating `MetricGaussianKL` again. For the
-    true probability distribution the standard parametrization is assumed.
-    The samples of this class are distributed among MPI tasks.
-