Repository: wtbarnes/fiasco

Issue:
numpy deprecation warning
Should probably fix this:

```python
/Users/dstansby/github/fiasco/fiasco/util/tools.py:175: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray
  nots = np.array(list(map(splrep_szero, x, y)))
```

Discussion:

Patch:
diff --git a/fiasco/collections.py b/fiasco/collections.py
index a74e11d..74ebcc6 100644
--- a/fiasco/collections.py
+++ b/fiasco/collections.py
@@ -41,11 +41,11 @@ class IonCollection:
             raise ValueError('Temperatures for all ions in collection must be the same.')
 
     def __getitem__(self, value):
-        ions = self._ion_list[value]
-        if isinstance(ions, list):
-            return IonCollection(*ions)
-        else:
+        ions = np.array(self._ion_list)[value]
+        if isinstance(ions, fiasco.Ion):
             return ions
+        else:
+            return IonCollection(*ions)
 
     def __contains__(self, value):
         if isinstance(value, (str, tuple)):
@@ -252,6 +252,7 @@ class IonCollection:
         """
         Calculate radiative loss curve which includes multiple ions
         """
+        density = np.atleast_1d(density)
         rad_loss = u.Quantity(np.zeros(self.temperature.shape + density.shape), 'erg cm^3 s^-1')
         for ion in self:
             try:
diff --git a/fiasco/elements.py b/fiasco/elements.py
index ee365b6..cf9b87f 100644
--- a/fiasco/elements.py
+++ b/fiasco/elements.py
@@ -9,6 +9,8 @@ from functools import cached_property
 
 import fiasco
 
+from fiasco.util import parse_ion_name
+
 __all__ = ['Element']
 
 
@@ -38,7 +40,7 @@ class Element(fiasco.IonCollection):
         Z = plasmapy.particles.atomic_number(element_name)
         ion_list = []
         for i in range(Z + 1):
-            ion = fiasco.Ion(f'{Z} {i+1}', temperature, **kwargs)
+            ion = fiasco.Ion((Z, i+1), temperature, **kwargs)
             ion_list.append(ion)
 
         super().__init__(*ion_list)
@@ -109,12 +111,9 @@ class Element(fiasco.IonCollection):
         return u.Quantity(ioneq)
 
     def __getitem__(self, value):
-        if type(value) is str:
-            el, ion = value.split()
-            if '+' in ion:
-                value = int(ion.strip('+'))
-            else:
-                value = int(ion) - 1
+        if isinstance(value, (str, tuple)):
+            _, value = parse_ion_name(value)
+            value -= 1
         return super().__getitem__(value)
 
     def __repr__(self):
diff --git a/fiasco/ions.py b/fiasco/ions.py
index 13078fb..208319e 100644
--- a/fiasco/ions.py
+++ b/fiasco/ions.py
@@ -147,6 +147,17 @@ Using Datasets:
                            self.temperature,
                            **self._instance_kwargs)
 
+    def previous_ion(self):
+        """
+        Return an `~fiasco.Ion` instance with the next lowest ionization stage.
+
+        For example, if the current instance is Fe XII (+11), this method returns
+        an instance of Fe XI (+10). All other input arguments remain the same.
+        """
+        return type(self)((self.atomic_number, self.ionization_stage-1),
+                           self.temperature,
+                           **self._instance_kwargs)
+
     @property
     @needs_dataset('elvlc', 'wgfa')
     def transitions(self):
@@ -426,9 +437,7 @@ Using Datasets:
         if include_protons:
             lower_level_p = self._psplups['lower_level']
             upper_level_p = self._psplups['upper_level']
-            pe_ratio = proton_electron_ratio(self.temperature,
-                                             **self._dset_names,
-                                             hdf5_dbase_root=self.hdf5_dbase_root)
+            pe_ratio = proton_electron_ratio(self.temperature, **self._instance_kwargs)
             proton_density = np.outer(pe_ratio, density)[:, :, np.newaxis]
             ex_rate_p = self.proton_collision_excitation_rate
             dex_rate_p = self.proton_collision_deexcitation_rate
diff --git a/fiasco/util/setup_db.py b/fiasco/util/setup_db.py
index 150a130..17e8ac7 100644
--- a/fiasco/util/setup_db.py
+++ b/fiasco/util/setup_db.py
@@ -15,12 +15,15 @@ import fiasco.io
 
 from fiasco.io import DataIndexer
 from fiasco.util.exceptions import MissingASCIIFileError, UnsupportedVersionError
-from fiasco.util.util import get_masterlist, query_yes_no
+from fiasco.util.util import get_chianti_catalog, query_yes_no
 
 FIASCO_HOME = pathlib.Path.home() / '.fiasco'
 CHIANTI_URL = 'http://download.chiantidatabase.org/CHIANTI_v{version}_database.tar.gz'
 # List in order (oldest to newest) the supported versions of the database
 SUPPORTED_VERSIONS = [
+    '8.0',
+    '8.0.2',
+    '8.0.6',
     '8.0.7',
 ]
 LATEST_VERSION = SUPPORTED_VERSIONS[-1]
@@ -132,7 +135,7 @@ def build_hdf5_dbase(ascii_dbase_root, hdf5_dbase_root, files=None):
 
     if files is None:
         files = []
-        tmp = get_masterlist(ascii_dbase_root)
+        tmp = get_chianti_catalog(ascii_dbase_root)
         for k in tmp:
             files += tmp[k]
     with ProgressBar(len(files)) as progress:
diff --git a/fiasco/util/tools.py b/fiasco/util/tools.py
index 7c5261f..264ecb6 100644
--- a/fiasco/util/tools.py
+++ b/fiasco/util/tools.py
@@ -152,9 +152,22 @@ def burgess_tully_descale(x, y, energy_ratio, c, scaling_type):
     # dtype (denoted by 'O') appear to be 1D, but should not be cast to 2D
     # as this will actually add an extra dimension and throw off the function
     # mapping
-    x = np.asarray(x)
-    y = np.asarray(y)
-    if x.dtype != np.dtype('O'):
+    # NOTE: We need to first work out whether the array is ragged or not in order
+    # to set the dtype of the resulting array. Not doing so is now deprecated in
+    # numpy. See https://github.com/wtbarnes/fiasco/issues/120
+    is_ragged = False
+    if isinstance(x, list):
+        if not all([_x.shape[0] == x[0].shape[0] for _x in x]):
+            is_ragged = True
+    elif isinstance(x, np.ndarray):
+        if x.dtype == np.dtype('O'):
+            is_ragged = True
+    else:
+        raise TypeError(f'x has unsupported type {type(x)}')
+    if is_ragged:
+        x = np.asarray(x, dtype='O')
+        y = np.asarray(y, dtype='O')
+    else:
         x = np.atleast_2d(x)
         y = np.atleast_2d(y)
     energy_ratio = np.atleast_2d(u.Quantity(energy_ratio).to_value(u.dimensionless_unscaled))
@@ -163,9 +176,9 @@ def burgess_tully_descale(x, y, energy_ratio, c, scaling_type):
     out = np.zeros(energy_ratio.shape)
     xnew = np.zeros(energy_ratio.shape)
 
-    for type in np.unique(scaling_type):
-        idxs = scaling_type == type
-        xnew[idxs, :] = _xnew(energy_ratio[idxs, :], c[idxs], type).T
+    for s_type in np.unique(scaling_type):
+        idxs = scaling_type == s_type
+        xnew[idxs, :] = _xnew(energy_ratio[idxs, :], c[idxs], s_type).T
 
     # Use list(map()) here to allow varying shaped inputs for x, y
     splrep_szero = partial(splrep, s=0)
@@ -173,17 +186,17 @@ def burgess_tully_descale(x, y, energy_ratio, c, scaling_type):
     splev_derzero = partial(splev, der=0)
     out = np.array(list(map(splev_derzero, xnew, nots)), dtype=object)
 
-    for type in np.unique(scaling_type):
-        idxs = scaling_type == type
-        if type == 1:
+    for s_type in np.unique(scaling_type):
+        idxs = scaling_type == s_type
+        if s_type == 1:
             out[idxs, ...] *= np.log(energy_ratio[idxs, ...] + np.e)
-        elif type == 3:
+        elif s_type == 3:
             out[idxs, ...] /= (energy_ratio[idxs, ...] + 1.0)
-        elif type == 4:
+        elif s_type == 4:
             out[idxs, ...] *= np.log(energy_ratio[idxs, ...].T + c[idxs]).T
-        elif type == 5:
+        elif s_type == 5:
             out[idxs, ...] /= energy_ratio[idxs, ...]
-        elif type == 6:
+        elif s_type == 6:
             out[idxs, ...] = 10**out[idxs]
 
     return out
diff --git a/fiasco/util/util.py b/fiasco/util/util.py
index 337cc1d..cec4bf2 100644
--- a/fiasco/util/util.py
+++ b/fiasco/util/util.py
@@ -12,7 +12,7 @@ from plasmapy.utils import roman
 FIASCO_HOME = pathlib.Path.home() / '.fiasco'
 FIASCO_RC = FIASCO_HOME / 'fiascorc'
 
-__all__ = ['setup_paths', 'get_masterlist', 'parse_ion_name']
+__all__ = ['setup_paths', 'get_chianti_catalog', 'parse_ion_name']
 
 
 def parse_ion_name(ion_name):
@@ -69,12 +69,26 @@ def setup_paths():
     return paths
 
 
-def get_masterlist(ascii_dbase_root):
+def get_chianti_catalog(ascii_dbase_root):
     """
+    Return a dictionary of all CHIANTI data files, separated by category.
+
     Parse CHIANTI filetree and return list of all files, separated by category. This will be only
     be useful when dealing with the raw ASCII data.
-    """
 
+    Parameters
+    ----------
+    ascii_dbase_root: path-like
+        Path to the top of the CHIANTI filetree
+
+    Returns
+    -------
+    : `dict`
+        All CHIANTI files, separated by category. The resulting dictionary should have the
+        following keys: 'abundance_files', 'ioneq_files', 'ip_files', 'continuum_files',
+        'ion_files'.
+    """
+    ascii_dbase_root = pathlib.Path(ascii_dbase_root)
     # TODO: Replace usage with pathlib, noting that pathlib does not
     # have a direct equivalent to os.walk
 


Test patch:
diff --git a/fiasco/tests/test_collections.py b/fiasco/tests/test_collections.py
index b180a0a..7753aa3 100644
--- a/fiasco/tests/test_collections.py
+++ b/fiasco/tests/test_collections.py
@@ -7,7 +7,8 @@ import pytest
 
 import fiasco
 
-temperature = np.logspace(5, 8, 100)*u.K
+temperature = np.logspace(4, 8, 100)*u.K
+wavelength = np.arange(1, 500, 1) * u.Angstrom
 
 
 @pytest.fixture
@@ -36,6 +37,14 @@ def collection(hdf5_dbase_root):
                                 fiasco.Ion('He 2', temperature, hdf5_dbase_root=hdf5_dbase_root))
 
 
+@pytest.fixture
+def another_collection(hdf5_dbase_root):
+    # This collection is for the continuum tests because we already have the needed
+    # files.
+    fe_5 = fiasco.Ion('Fe 5', temperature, hdf5_dbase_root=hdf5_dbase_root)
+    return fe_5 + fe_5.next_ion()
+
+
 def test_create_collection_from_ions(ion, another_ion):
     assert isinstance(ion + another_ion, fiasco.IonCollection)
     assert isinstance(fiasco.IonCollection(ion, another_ion), fiasco.IonCollection)
@@ -48,6 +57,8 @@ def test_create_collection_from_elements(element, another_element):
 
 def test_create_collection_from_mixture(ion, element, collection):
     assert isinstance(ion + element + collection, fiasco.IonCollection)
+    # This tests the reflected operator
+    assert isinstance(collection + ion, fiasco.IonCollection)
     assert isinstance(fiasco.IonCollection(ion, element, collection), fiasco.IonCollection)
 
 
@@ -64,6 +75,7 @@ def test_getitem(ion, another_ion, element, another_element):
     assert isinstance(collection[:2], fiasco.IonCollection)
     for i in collection:
         assert isinstance(i, fiasco.Ion)
+    assert isinstance(collection[[1,2]], fiasco.IonCollection)
 
 
 def test_contains(collection, hdf5_dbase_root):
@@ -74,6 +86,28 @@ def test_contains(collection, hdf5_dbase_root):
     assert ion in collection
 
 
+def test_length(collection):
+    assert len(collection) == len(collection._ion_list)
+
+
+def test_free_free(another_collection):
+    ff = another_collection.free_free(wavelength)
+    assert ff.shape == temperature.shape + wavelength.shape
+    assert u.allclose(ff[50, 50], 3.19877384e-35 * u.Unit('erg cm3 s-1 Angstrom-1'))
+
+
+def test_free_bound(another_collection):
+    fb = another_collection.free_bound(wavelength)
+    assert fb.shape == temperature.shape + wavelength.shape
+    assert u.allclose(fb[50, 50], 3.2653516e-29 * u.Unit('erg cm3 s-1 Angstrom-1'))
+
+
+def test_radiative_los(collection):
+    rl = collection.radiative_loss(1e9*u.cm**(-3))
+    # This value has not been checked for correctness
+    assert u.allclose(rl[0,0], 3.2389535764824023e-24*u.Unit('erg cm3 s-1'))
+
+
 def test_spectrum(hdf5_dbase_root):
     i1 = fiasco.Ion('H 1', 1 * u.MK, hdf5_dbase_root=hdf5_dbase_root)
     i2 = fiasco.Ion('Fe 5', 1 * u.MK, hdf5_dbase_root=hdf5_dbase_root)
diff --git a/fiasco/tests/test_element.py b/fiasco/tests/test_element.py
index 2ac4e33..ac3d4d2 100644
--- a/fiasco/tests/test_element.py
+++ b/fiasco/tests/test_element.py
@@ -54,7 +54,12 @@ def test_create_element_number(symbol, hdf5_dbase_root):
 
 def test_getitem_ion_name(element):
     assert element['H 1'] == element[0]
-    assert element['H 2'] == element[-1]
+    assert element['H 1'].ionization_stage == 1
+    assert element['H 1'].atomic_number == 1
+    assert element['H 2'] == element[1]
+    assert element['H 2'].ionization_stage == 2
+    assert element['H 2'].atomic_number == 1
+    assert element['H +1'] == element[1]
 
 
 def test_create_element_without_units_raises_units_error(hdf5_dbase_root):
@@ -67,6 +72,14 @@ def test_create_element_with_wrong_units_raises_unit_conversion_error(hdf5_dbase
         fiasco.Element('H', temperature.value*u.s, hdf5_dbase_root=hdf5_dbase_root)
 
 
-def test_equilibrium_ionization(element):
-    ioneq = element.equilibrium_ionization
-    assert ioneq.shape == element.temperature.shape + (element.atomic_number + 1,)
+def test_equilibrium_ionization(hdf5_dbase_root):
+    # NOTE: Using an element with Z>1 so that we are testing the full rate matrix
+    # computation. Using C here because we are already using this for the gallery
+    carbon = fiasco.Element('C', temperature, hdf5_dbase_root=hdf5_dbase_root)
+    ioneq = carbon.equilibrium_ionization
+    assert ioneq.shape == carbon.temperature.shape + (carbon.atomic_number + 1,)
+    assert ioneq[33, 5] == u.Quantity(0.5787345345914312)
+
+
+def test_element_repr(element):
+    assert element.element_name in element.__repr__()
diff --git a/fiasco/tests/test_ion.py b/fiasco/tests/test_ion.py
index 444bb5b..e747f49 100644
--- a/fiasco/tests/test_ion.py
+++ b/fiasco/tests/test_ion.py
@@ -263,3 +263,15 @@ def test_repr_no_levels(hdf5_dbase_root):
     no energy level or transition information is available.
     """
     assert fiasco.Ion('Fe 1', temperature, hdf5_dbase_root=hdf5_dbase_root).__repr__
+
+
+def test_next_ion(ion):
+    next_ion = ion.next_ion()
+    assert next_ion.ionization_stage == ion.ionization_stage + 1
+    assert next_ion.atomic_number == ion.atomic_number
+
+
+def test_previous_ion(ion):
+    prev_ion = ion.previous_ion()
+    assert prev_ion.ionization_stage == ion.ionization_stage - 1
+    assert prev_ion.atomic_number == ion.atomic_number
diff --git a/fiasco/util/tests/test_tools.py b/fiasco/util/tests/test_tools.py
index 14afb61..7b83f70 100644
--- a/fiasco/util/tests/test_tools.py
+++ b/fiasco/util/tests/test_tools.py
@@ -1,23 +1,23 @@
 import numpy as np
+import pytest
 
 from fiasco.util import burgess_tully_descale
 
 
-def test_burgess_tully_one_array():
-    x = np.linspace(0, 1, 10)
-    y = x**1.6 + 10
-    energy_ratio = np.ones((1,) + x.shape)
-    c = 1
-    for scaling_type in range(1, 7):
-        ups = burgess_tully_descale(x, y, energy_ratio, c, scaling_type)
-        assert ups.shape == energy_ratio.shape
-
-
-def test_burgess_tully_2d_array():
-    x = np.tile(np.linspace(0, 1, 10), (2, 1))
-    y = x**1.6 + 10
-    energy_ratio = np.ones(x.shape[:1] + (10,))
-    c = 1*np.ones(x.shape[:1])
+@pytest.mark.parametrize('x', [
+    np.linspace(0, 1, 10).reshape(1,10),  # one array
+    np.tile(np.linspace(0, 1, 10), (2, 1)),  # 2D array
+    [np.linspace(0,1,10), np.linspace(0,1,10)], # list with equal lengths
+    np.array([np.linspace(0, 1, 5), np.linspace(0, 1, 9)], dtype=object),  # staggered array
+    [np.linspace(0,1,5), np.linspace(0,1,9)], # list with unequal lengths
+])
+def test_burgess_tully(x):
+    if isinstance(x, list):
+        y = [_x**1.6 + 10 for _x in x]
+    else:
+        y = x**1.6 + 10
+    energy_ratio = np.ones((len(x), 10))
+    c = 1*np.ones((len(x),))
     for scaling_type in range(1, 7):
         st = scaling_type*np.ones(c.shape)
         ups = burgess_tully_descale(x, y, energy_ratio, c, st)
@@ -32,14 +32,3 @@ def test_burgess_tully_different_scaling_types():
     scaling_type = np.arange(1, 7)
     ups = burgess_tully_descale(x, y, energy_ratio, c, scaling_type)
     assert ups.shape == energy_ratio.shape
-
-
-def test_burgess_tully_staggered_array():
-    x = np.array([np.linspace(0, 1, 5), np.linspace(0, 1, 9)], dtype=object)
-    y = x**1.6 + 10
-    energy_ratio = np.ones(x.shape[:1] + (10,))
-    c = 1*np.ones(x.shape)
-    for scaling_type in range(1, 7):
-        st = scaling_type*np.ones(c.shape)
-        ups = burgess_tully_descale(x, y, energy_ratio, c, st)
-        assert ups.shape == energy_ratio.shape
diff --git a/fiasco/util/tests/test_util.py b/fiasco/util/tests/test_util.py
index 8e60c37..4fffc77 100644
--- a/fiasco/util/tests/test_util.py
+++ b/fiasco/util/tests/test_util.py
@@ -3,7 +3,7 @@ Tests for util functions
 """
 import pytest
 
-from fiasco.util import parse_ion_name
+from fiasco.util import get_chianti_catalog, parse_ion_name
 
 
 @pytest.mark.parametrize('ion_name', [
@@ -23,3 +23,15 @@ def test_parse_ion_name(ion_name):
     element, ion = parse_ion_name(ion_name)
     assert element == 26
     assert ion == 21
+
+
+def test_get_chianti_catalog(ascii_dbase_root):
+    catalog = get_chianti_catalog(ascii_dbase_root)
+    keys = ['abundance_files',
+            'ioneq_files',
+            'ip_files',
+            'continuum_files',
+            'ion_files']
+    for k in catalog:
+        assert k in catalog
+        assert isinstance(catalog[k], list)