Fixes to Python module

hyperion/densities/ambient_medium.py

@@ -56,6 +56,9 @@ def __init__(self, rho=None, rmin=None, rmax=None, subtract=[]):
         self.rmax = rmax
         self.subtract = subtract
 
+        # Central star
+        self.star = None
+
         # Dust
         self.dust = None
 

hyperion/densities/power_law_envelope.py

@@ -118,9 +118,8 @@ def mass(self):
             return None
         else:
             self._check_all_set()
-            alpha = 3. + self.power
-            mass = self.rho_0 / alpha * \
-                (4. * pi * (self.rmax ** alpha - self.rmin ** alpha) / self.r_0 ** self.power)
+            mass = self.rho_0 * \
+                (4. * pi * integrate_powerlaw(self.rmin, self.rmax, 2. + self.power) / self.r_0 ** self.power)
             return mass
 
     @mass.setter
@@ -141,9 +140,8 @@ def rho_0(self):
             return None
         else:
             self._check_all_set()
-            alpha = 3. + self.power
-            rho_0 = self.mass * alpha / \
-                (4. * pi * (self.rmax ** alpha - self.rmin ** alpha) / self.r_0 ** self.power)
+            rho_0 = self.mass / \
+                (4. * pi * integrate_powerlaw(self.rmin, self.rmax, 2. + self.power) / self.r_0 ** self.power)
             return rho_0
 
     @rho_0.setter

hyperion/densities/tests/test_densities.py

@@ -888,3 +888,26 @@ def test_alpha_disk_zero_discretized_mass():
     with pytest.raises(Exception) as exc:
         e.density(g)
     assert "zero" in exc.value.args[0]
+
+
+def test_power_law_envelope_power_minus_three():
+    # power = -3 gives alpha = 3 + power = 0, which used to divide by zero in
+    # the mass <-> rho_0 conversion; it should now round-trip cleanly.
+    e = PowerLawEnvelope()
+    e.rmin = 1.
+    e.rmax = 10.
+    e.r_0 = 5.
+    e.power = -3.
+    e.rho_0 = 2.
+    m = e.mass
+    assert np.isfinite(m) and m > 0.
+    e.mass = m
+    np.testing.assert_allclose(e.rho_0, 2.)
+
+
+def test_ambient_medium_star_attribute():
+    # AmbientMedium.star must be declared before the instance is frozen,
+    # otherwise assigning it raises AttributeError.
+    a = AmbientMedium()
+    a.star = Star()
+    assert a.star is not None

hyperion/grid/amr_grid.py

@@ -15,9 +15,9 @@
 from .grid_helpers import single_grid_dims
 
 
-def zero_density(grid, xmin=-np.inf, xmax=np.inf, ymin=-np.inf, ymax=np.inf, zmin=np.inf, zmax=np.inf):
-    for ilevel, level in enumerate(grid.levels):
-        for igrid, grid in enumerate(level.grids):
+def zero_density(amr_grid, xmin=-np.inf, xmax=np.inf, ymin=-np.inf, ymax=np.inf, zmin=-np.inf, zmax=np.inf):
+    for level in amr_grid.levels:
+        for grid in level.grids:
             wx = np.linspace(grid.xmin, grid.xmax, grid.nx + 1)
             wy = np.linspace(grid.ymin, grid.ymax, grid.ny + 1)
             wz = np.linspace(grid.zmin, grid.zmax, grid.nz + 1)
@@ -26,8 +26,13 @@ def zero_density(grid, xmin=-np.inf, xmax=np.inf, ymin=-np.inf, ymax=np.inf, zmi
             z = 0.5 * (wz[:-1] + wz[1:])
             gx, gy, gz = meshgrid_nd(x, y, z)
             reset = (gx < xmin) | (gx > xmax) | (gy < ymin) | (gy > ymax) | (gz < zmin) | (gz > zmax)
-            grid.data[reset] = 0.
-    return grid
+            for quantity in grid.quantities:
+                if isinstance(grid.quantities[quantity], list):
+                    for array in grid.quantities[quantity]:
+                        array[reset] = 0.
+                else:
+                    grid.quantities[quantity][reset] = 0.
+    return amr_grid
 
 
 class Grid(FreezableClass):

hyperion/grid/tests/test_voronoi.py

@@ -147,3 +147,19 @@ def test_setitem_from_view():
     assert g2.xmin == g1.xmin and g2.xmax == g1.xmax
     assert g2.ymin == g1.ymin and g2.ymax == g1.ymax
     assert g2.zmin == g1.zmin and g2.zmax == g1.zmax
+
+
+def test_evaluate_function_average_constant():
+    # For a constant function every (non-empty) cell average must equal that
+    # constant. The old code divided the last cell's sum by a decremented
+    # sample count, inflating that one cell's average away from the constant.
+    np.random.seed(9876)
+    N = 12
+    x = np.random.random(N)
+    y = np.random.random(N)
+    z = np.random.random(N)
+
+    g = VoronoiGrid(x, y, z)
+    avg = g.evaluate_function_average(lambda x, y, z: 3.5 * np.ones_like(x),
+                                      n_samples=2000000, min_cell_samples=5)
+    np.testing.assert_allclose(avg, 3.5, rtol=1e-6)

hyperion/grid/tests/test_zero_density.py

@@ -0,0 +1,35 @@
+import numpy as np
+
+from ..amr_grid import AMRGrid, zero_density
+
+
+def _single_grid_amr():
+    amr = AMRGrid()
+    level = amr.add_level()
+    grid = level.add_grid()
+    grid.xmin, grid.xmax = -1., 1.
+    grid.ymin, grid.ymax = -1., 1.
+    grid.zmin, grid.zmax = -1., 1.
+    grid.nx, grid.ny, grid.nz = 4, 4, 4
+    grid.quantities['density'] = np.ones((4, 4, 4))
+    return amr
+
+
+def test_zero_density_default_keeps_density():
+    # Regression: zmin used to default to +inf, so the reset mask covered
+    # every cell and the whole grid was zeroed. With no limits given, nothing
+    # should be zeroed.
+    amr = _single_grid_amr()
+    out = zero_density(amr)
+    assert out is amr                                  # returns the grid passed in
+    assert np.all(out.levels[0].grids[0].quantities['density'] == 1.)
+
+
+def test_zero_density_zeroes_outside_box():
+    # Only cells with z > 0 should be zeroed.
+    amr = _single_grid_amr()
+    zero_density(amr, zmin=-np.inf, zmax=0.)
+    dens = amr.levels[0].grids[0].quantities['density']
+    # cell z-centres are at -0.75, -0.25, 0.25, 0.75 (first axis is z)
+    assert np.all(dens[2:] == 0.)                      # z > 0 zeroed
+    assert np.all(dens[:2] == 1.)                      # z < 0 kept

hyperion/grid/voronoi_grid.py

@@ -203,20 +203,21 @@ def evaluate_function_average(self, func, n_samples=None, min_cell_samples=None)
         else:
             single = False
 
-        # reduceat does not recognize values of indices that are the size of
-        # the array so we need to subtract one here for those values.
-        idx = self._samples_idx.copy()
-        idx[idx == len(values_all[0])] -= 1
+        # reduceat does not accept indices equal to the size of the array
+        # (which occur for empty cells at the end of the grid) and for two
+        # consecutive equal index values (empty cells) it returns the value
+        # at that index rather than 0, so we only pass the start indices of
+        # the non-empty cells and explicitly set the average to 0 for empty
+        # cells.
+        counts = np.diff(self._samples_idx)
+        non_empty = counts > 0
+        starts = self._samples_idx[:-1][non_empty]
 
         averages_all = []
 
         for values in values_all:
-            averages = np.add.reduceat(values, idx[:-1]) / np.diff(idx)
-
-            # reduceat doesn't do what we want for two consecutive equal
-            # index values - it will take the value at that index rather than
-            # return 0 (which is what would be expected for i:i slicing)
-            averages[np.isinf(averages)] = 0.
+            averages = np.zeros(len(counts))
+            averages[non_empty] = np.add.reduceat(values, starts) / counts[non_empty]
             averages_all.append(averages)
 
         if single:

hyperion/importers/_discretize_sph.c

@@ -262,34 +262,49 @@ static PyObject *_get_positions_widths(PyObject *self, PyObject *args)
     PyObject *yc_array = PyArray_SimpleNew(1, dims, NPY_DOUBLE);
     if (yc_array == NULL) {
         PyErr_SetString(PyExc_TypeError, "Couldn't build output array");
+        Py_XDECREF(xc_array);
         Py_XDECREF(refined_array);
         return NULL;
     }
 
     PyObject *zc_array = PyArray_SimpleNew(1, dims, NPY_DOUBLE);
     if (zc_array == NULL) {
         PyErr_SetString(PyExc_TypeError, "Couldn't build output array");
+        Py_XDECREF(xc_array);
+        Py_XDECREF(yc_array);
         Py_XDECREF(refined_array);
         return NULL;
     }
 
     PyObject *xw_array = PyArray_SimpleNew(1, dims, NPY_DOUBLE);
     if (xw_array == NULL) {
         PyErr_SetString(PyExc_TypeError, "Couldn't build output array");
+        Py_XDECREF(xc_array);
+        Py_XDECREF(yc_array);
+        Py_XDECREF(zc_array);
         Py_XDECREF(refined_array);
         return NULL;
     }
 
     PyObject *yw_array = PyArray_SimpleNew(1, dims, NPY_DOUBLE);
     if (yw_array == NULL) {
         PyErr_SetString(PyExc_TypeError, "Couldn't build output array");
+        Py_XDECREF(xc_array);
+        Py_XDECREF(yc_array);
+        Py_XDECREF(zc_array);
+        Py_XDECREF(xw_array);
         Py_XDECREF(refined_array);
         return NULL;
     }
 
     PyObject *zw_array = PyArray_SimpleNew(1, dims, NPY_DOUBLE);
     if (zw_array == NULL) {
         PyErr_SetString(PyExc_TypeError, "Couldn't build output array");
+        Py_XDECREF(xc_array);
+        Py_XDECREF(yc_array);
+        Py_XDECREF(zc_array);
+        Py_XDECREF(xw_array);
+        Py_XDECREF(yw_array);
         Py_XDECREF(refined_array);
         return NULL;
     }
@@ -309,11 +324,22 @@ static PyObject *_get_positions_widths(PyObject *self, PyObject *args)
 
     if(i != ncells - 1) {
         PyErr_SetString(PyExc_TypeError, "An error occurred when retrieving the cell properties");
+        Py_XDECREF(xc_array);
+        Py_XDECREF(yc_array);
+        Py_XDECREF(zc_array);
+        Py_XDECREF(xw_array);
+        Py_XDECREF(yw_array);
+        Py_XDECREF(zw_array);
         Py_XDECREF(refined_array);
         return NULL;
     }
 
-    return Py_BuildValue("OOOOOO", xc_array, yc_array, zc_array, xw_array, yw_array, zw_array);
+    /* Clean up. */
+    Py_XDECREF(refined_array);
+
+    /* Use the "N" format code, which steals the references to the output
+       arrays instead of adding new ones. */
+    return Py_BuildValue("NNNNNN", xc_array, yc_array, zc_array, xw_array, yw_array, zw_array);
 
 }
 

hyperion/importers/tests/test_sph.py

@@ -68,3 +68,27 @@ def get_volumes(current_i, current_volume, refined):
     volumes_ref.insert(0, 6 * 5 * 4 * 8.)
 
     assert_allclose(volumes, volumes_ref)
+
+
+def test_get_positions_widths_no_reference_leak():
+    # Regression: _get_positions_widths leaked references to its input array
+    # (never DECREF'd) and to its six output arrays (returned with "O" instead
+    # of "N"). Check both are now reference-clean.
+    import sys
+    import numpy as np
+    from .._discretize_sph import _get_positions_widths
+
+    refined = np.array([True] + [False] * 8, dtype=bool)
+
+    rc0 = sys.getrefcount(refined)
+    for _ in range(20):
+        out = _get_positions_widths(refined, 0., 0., 0., 1., 1., 1.)
+        del out
+    assert sys.getrefcount(refined) == rc0            # input not leaked
+
+    out = _get_positions_widths(refined, 0., 0., 0., 1., 1., 1.)
+    arr = out[0]
+    del out
+    # arr is now referenced only by `arr` (+ the temporary getrefcount arg);
+    # the "O" leak would leave an extra dangling reference.
+    assert sys.getrefcount(arr) == 2

hyperion/util/functions.py

@@ -80,7 +80,6 @@ class FreezableClass(object):
 
     _frozen = False
     _final = False
-    _attributes = []
 
     def __init__(self):
         super(FreezableClass, self).__init__()
@@ -100,9 +99,11 @@ def isfinal(self):
     def __setattr__(self, key, value):
         if self._final:
             raise Exception("Attribute %s can no longer be changed" % key)
+        if '_attributes' not in self.__dict__:
+            object.__setattr__(self, '_attributes', set())
         if self._frozen and not key in self._attributes:
             raise AttributeError("Attribute %s does not exist" % key)
-        self._attributes.append(key)
+        self._attributes.add(key)
         object.__setattr__(self, key, value)
 
 

hyperion/util/tests/test_functions.py

@@ -1,3 +1,4 @@
+import pytest
 import numpy as np
 
 from ..functions import B_nu, dB_nu_dT
@@ -40,3 +41,28 @@ def test_db_nu_dt():
 
         # Check that the two are the same
         np.testing.assert_allclose(db, db_num, rtol=1.e-2)
+
+
+def test_freezable_attributes_per_instance():
+    # Regression: FreezableClass used to store _attributes on the class, so
+    # attributes registered on one instance leaked into every other instance
+    # (and grew an unbounded shared list). They must be per-instance.
+    from ..functions import FreezableClass
+
+    class Foo(FreezableClass):
+        def __init__(self):
+            FreezableClass.__init__(self)
+
+    a = Foo()
+    a.x = 1        # registers 'x' on a only
+    a._freeze()
+
+    b = Foo()
+    b._freeze()
+    # 'x' belongs to a, not b; with the shared-list bug b would allow it
+    with pytest.raises(AttributeError):
+        b.x = 2
+
+    # a can still change its own registered attribute
+    a.x = 5
+    assert a.x == 5