[FIX] Bug fix in num_dens calculation

astro_plasma/core/ionization.py

@@ -200,7 +200,7 @@ def interpolate_ion_frac(
             if _is_multiple:
                 return fracIon.flatten().reshape((*self._input_shape, fracIon.shape[-1]))
             else:
-                return fracIon
+                return fracIon.flatten()
 
         _element, _ion = parse_atomic_ion_no(element, ion)
 
@@ -229,7 +229,7 @@ def interpolate_ion_frac(
             else:
                 return fracIon.flatten().reshape((*self._input_shape, _element + 1))
         else:
-            fracIon = self._interpolate_ion_frac_all(nH, temperature, metallicity, redshift, mode)[slice_start:slice_stop]
+            fracIon = self._interpolate_ion_frac_all(nH, temperature, metallicity, redshift, mode).flatten()[slice_start:slice_stop]
             # Array starts from 0 but _ion from 1
             return fracIon[_ion - 1] if _ion is not None else fracIon  # This is in log10
 
@@ -302,7 +302,7 @@ def interpolate_num_dens(
             for element in range(30):
                 for ion in range(element + 2):
                     slices.append(slice(ion_count, ion_count + 1))
-                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon[ion_count]
+                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon.flatten()[ion_count]
                     slices = slices[:-1]
                     if element + 1 == 1:  # H
                         ndens += (ion + 1) * (Xp(metallicity) / X_solar) * abn[element] * _ion_fraction * nH
@@ -319,7 +319,7 @@ def interpolate_num_dens(
             for element in range(30):
                 for ion in range(element + 2):
                     slices.append(slice(ion_count, ion_count + 1))
-                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon[ion_count]
+                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon.flatten()[ion_count]
                     slices = slices[:-1]
                     if element + 1 == 1:  # H
                         ne += ion * (Xp(metallicity) / X_solar) * nH * abn[element] * _ion_fraction
@@ -336,7 +336,7 @@ def interpolate_num_dens(
             for element in range(30):
                 for ion in range(element + 2):
                     slices.append(slice(ion_count, ion_count + 1))
-                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon[ion_count]
+                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon.flatten()[ion_count]
                     slices = slices[:-1]
                     if ion == 0:
                         if element + 1 == 1:  # H
@@ -355,7 +355,7 @@ def interpolate_num_dens(
                 ion_count += 1  # neglects the neutral species
                 for ion in range(1, element + 2):
                     slices.append(slice(ion_count, ion_count + 1))
-                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon[ion_count]
+                    _ion_fraction = fracIon[tuple(slices)].flatten().reshape(self._input_shape) if _is_multiple else fracIon.flatten()[ion_count]
                     slices = slices[:-1]
                     if element + 1 == 1:  # H
                         nion += (Xp(metallicity) / X_solar) * nH * abn[element] * _ion_fraction

tests/test_astro_palsma_basic.py

@@ -86,6 +86,19 @@ def test_dimension():
     )  # This value is in log10
     assert np.array(frac).ndim == 0
 
+    frac = fIon(
+        nH=nH,
+        temperature=[
+            temperature,
+        ],
+        metallicity=metallicity,
+        redshift=redshift,
+        element=element,
+        ion=ion,
+        mode=mode,
+    )  # This value is in log10
+    assert np.array(frac).ndim == 0
+
     nH = np.power(10.0, np.random.uniform(low=-5.0, high=-2.0))  # Hydrogen number density in cm^-3
     temperature = np.power(10.0, np.linspace(3.8, 6.5, 5))  # Temperature of the plasma in kelvin
     metallicity = np.random.uniform(low=0.2, high=0.9)  # Metallicity of plasma with respect to solar

tests/test_astro_palsma_run.py

@@ -52,12 +52,13 @@ def test_num_dens():
 
     ne = num_dens(
         nH=nH,
-        temperature=temperature,
+        temperature=[temperature],
         metallicity=metallicity,
         redshift=redshift,
         mode=mode,
         part_type="electron",
     )
+    assert np.array(ne).ndim == 0
     ne_expected = 1.4109277149716788e-04
     assert np.isclose(ne, ne_expected)