Merge pull request #16 from mirochaj/meshgrid_indexing

providing `indexing` to numpy.meshgrid

CHANGELOG.rst

@@ -17,6 +17,9 @@ Changelog
   inverse transforms remain inverses by default, the phases are interpreted as
   having zero at the centre (for both transforms). See the phasing tutorial for
   more information.
+- Fixed transpose issue caused by default behavior of ``numpy.meshgrid``, which
+  led to broken correspondence between discrete sample of field and original
+  field. See [Issue #15].
 
 **Bugfixes**
 

CONTRIBUTORS.rst

@@ -5,4 +5,5 @@ Authors
 
 Comments, corrections and suggestions
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-* `Chris Jordan <https://github.com/cjordan>`_
+* `Chris Jordan <https://github.com/cjordan>`_
+* `Jordan Mirocha <https://github.com/mirochaj>`_

powerbox/powerbox.py

@@ -6,6 +6,7 @@
 over-writing the same methods as are over-written in :class:`LogNormalPowerBox`.
 """
 
+import warnings
 import numpy as np
 from . import dft
 from .tools import _magnitude_grid
@@ -230,7 +231,7 @@ def delta_x(self):
         return dk
 
     def create_discrete_sample(self, nbar, randomise_in_cell=True, min_at_zero=False,
-                               store_pos=False):
+                               store_pos=False, delta_x=None):
         r"""
         Assuming that the real-space signal represents an over-density with respect to some mean, create a sample
         of tracers of the underlying density distribution.
@@ -247,21 +248,40 @@ def create_discrete_sample(self, nbar, randomise_in_cell=True, min_at_zero=False
             origin.
         store_pos : bool, optional
             Whether to store the sample of tracers as an instance variable `tracer_positions`.
+        delta_x : numpy.ndarray
+            Field from which to draw discrete samples. This is likely the
+            output of a previous call to `delta_x()`, but could in principle be
+            any field. Note that if not supplied, the field will be generated
+            from scratch. As a result, unless the user has supplied a random seed
+            at initialization, the discrete samples will be a new realization of
+            a field with the specified power spectrum.
 
         Returns
         -------
         tracer_positions : float, array_like
             ``(n, d)``-array, with ``n`` the number of tracers and ``d`` the number of dimensions. Each row represents
             a single tracer's co-ordinates.
         """
-        dx = self.delta_x()
+
+        if delta_x is None:
+            if self.seed is None:
+                warnings.warn(
+                    "WARNING: should provide `seed` at initialization if one"
+                    " wants a correspondence between parent field and"
+                    " discrete samples."
+                )
+            dx = self.delta_x()
+        else:
+            dx = delta_x
+
         dx = (dx + 1) * self.dx ** self.dim * nbar
         n = dx
+
         self.n_per_cell = np.random.poisson(n)
 
         # Get all source positions
         args = [self.x] * self.dim
-        X = np.meshgrid(*args)
+        X = np.meshgrid(*args, indexing='ij')
 
         tracer_positions = np.array([x.flatten() for x in X]).T
         tracer_positions = tracer_positions.repeat(self.n_per_cell.flatten(), axis=0)

tests/test_discrete.py

@@ -8,19 +8,33 @@
 
 from powerbox import PowerBox, LogNormalPowerBox, get_power
 
-
 def test_discrete_power_gaussian():
-    pb = PowerBox(N=512,dim=2,boxlength=100., pk = lambda u : 0.1*u**-1.5, ensure_physical=True)
+    pb = PowerBox(N=512,dim=2,boxlength=100., pk = lambda u : 0.1*u**-1.5,
+        ensure_physical=True)
 
-    sample = pb.create_discrete_sample(nbar=1000.)
-    power, bins = get_power(sample,pb.boxlength,N=pb.N)
+    box = pb.delta_x()
 
+    sample = pb.create_discrete_sample(nbar=1000., delta_x=box)
+    power, bins = get_power(sample,pb.boxlength,N=pb.N)
 
     res =  np.mean(np.abs(power[50:-50] / (0.1*bins[50:-50]**-1.5) -1 ) )
 
-    print(res)
     assert res < 1e-1
 
+    # This re-grids the discrete sample into a box, basically to verify the
+    # indexing used by meshgrid within `create_discrete_sample`.
+    N = [pb.N] * pb.dim
+    L = [pb.boxlength] * pb.dim
+    edges = [np.linspace(-_L/2., _L/2., _n + 1) for _L, _n in zip(L, N)]
+    delta_samp = np.histogramdd(sample, bins=edges, weights=None)[0].astype("float")
+
+    # Check cross spectrum and assert a strong correlation
+    cross, bins = get_power(delta_samp,pb.boxlength,deltax2=box)
+    p2, bins = get_power(box, pb.boxlength)
+    corr = cross / np.sqrt(power) / np.sqrt(p2)
+    corr_bar = np.mean(corr[np.isfinite(corr)])
+    assert corr_bar > 10
+
 def test_discrete_power_lognormal():
     pb = LogNormalPowerBox(N=512, dim=2, boxlength=100., pk=lambda u: 0.1*u ** -1.5, ensure_physical=True)
 
@@ -29,5 +43,8 @@ def test_discrete_power_lognormal():
 
     res = np.mean(np.abs(power[50:-50]/(0.1*bins[50:-50] ** -1.5) - 1))
 
-    print(res)
-    assert res < 1e-1
+    assert res < 1e-1
+
+if __name__ == '__main__':
+    test_discrete_power_gaussian()
+    test_discrete_power_lognormal()