Modify ion.n_levels

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ ci:
 
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.12.7"
+    rev: "v0.14.4"
     hooks:
       - id: ruff
         args: ["--fix"]

.ruff.toml

@@ -0,0 +1,49 @@
+target-version = "py312"
+line-length = 110
+exclude=[
+  ".git,",
+  "__pycache__",
+  "build",
+  "fiasco/version.py",
+]
+show-fixes = true
+output-format = "full"
+
+[lint]
+select = [
+  "E",
+  "F",
+  "W",
+  "UP",
+  "PT",
+  #"RET",
+  #"TID",
+]
+extend-ignore = [
+  # pycodestyle (E, W)
+  "E501",  # LineTooLong                             # TODO! fix
+  "E741",  # Ambiguous variable name
+
+  # pytest (PT)
+  "PT001",  # Always use pytest.fixture()
+  "PT004",  # Fixtures which don't return anything should have leading _
+  "PT007",  # Parametrize should be lists of tuples  # TODO! fix
+  "PT011",  # Too broad exception assert             # TODO! fix
+  "PT023",  # Always use () on pytest decorators
+]
+
+[lint.per-file-ignores]
+# Part of configuration, not a package.
+"setup.py" = ["INP001"]
+"conftest.py" = ["INP001"]
+# implicit-namespace-package. The examples are not a package.
+"docs/*.py" = ["INP001"]
+# Module level imports do not need to be at the top of a file here
+"docs/conf.py" = ["E402"]
+"__init__.py" = ["E402", "F401", "F403"]
+"test_*.py" = ["B011", "D", "E402", "S101"]
+# Allow star import for the factory
+"fiasco/io/factory.py" = ["F403"]
+
+[lint.pydocstyle]
+convention = "numpy"

examples/idl_comparisons/continuum_comparison.py

@@ -104,7 +104,9 @@ def plot_idl_comparison(wavelength, temperature, result_fiasco, result_idl):
               'ionization_fraction': idl_result_twophoton['ionization_fraction']}
 all_ions = [fiasco.Ion(ion_name, idl_result_twophoton['temperature'], **ion_kwargs) for ion_name in fiasco.list_ions()]
 all_ions = fiasco.IonCollection(*all_ions)
-two_photon = all_ions.two_photon(idl_result_twophoton['wavelength'], idl_result_twophoton['density']).squeeze()
+two_photon = all_ions.two_photon(idl_result_twophoton['wavelength'],
+                                 idl_result_twophoton['density'],
+                                 use_two_ion_model=False).squeeze()
 plot_idl_comparison(idl_result_twophoton['wavelength'],
                     idl_result_twophoton['temperature'],
                     two_photon,

examples/user_guide/aia_response.py

@@ -42,7 +42,7 @@
 ############################################################
 # Next, we construct for the `~fiasco.Ion` object for Fe 18. Note
 # that by default we use the coronal abundances of :cite:t:`feldman_potential_1992`.
-temperature = 10.**(np.arange(4.5, 8, 0.05)) * u.K
+temperature = 10.**(np.arange(4.5, 8, 0.1)) * u.K
 fe18 = fiasco.Ion('Fe 18', temperature)
 
 ############################################################

examples/user_guide/ion_collection.py

@@ -25,7 +25,7 @@
 #
 # We will create `~fiasco.IonCollection` containing both O VI
 # and Fe XVIII
-temperature = np.geomspace(1e4,1e8,100) * u.K
+temperature = np.geomspace(10**5, 10**7.5, 30) * u.K
 fe18 = fiasco.Ion('Fe XVIII', temperature)
 o6 = fiasco.Ion('O VI', temperature)
 col = fiasco.IonCollection(fe18, o6)

examples/user_guide/plot_contribution_function.py

@@ -19,7 +19,7 @@
 # Specify the plasma properties; note that an `~fiasco.Ion` has to be created with a
 # range of temperatures, but the density is only used later in the contribution
 # function calculation.
-Te = np.geomspace(0.1, 100, 51) * u.MK
+Te = np.geomspace(0.1, 10, 26) * u.MK
 ne = 1e8 * u.cm**-3
 
 ###############################################################################

fiasco/ions.py

@@ -108,6 +108,9 @@ def __repr__(self):
     @cached_property
     @needs_dataset('elvlc')
     def levels(self):
+        """
+        Information for each energy level of the ion.
+        """
         return Levels(self._elvlc)
 
     def __getitem__(self, key):
@@ -159,12 +162,28 @@ def _has_dataset(self, dset_name):
     @property
     def n_levels(self):
         """
-        Number of energy levels in the CHIANTI model
+        Number of energy levels in the atomic model.
+
+        .. note:: It is possible this number will not match the number of levels
+                  in `~fiasco.Ion.levels`. The number of levels in a model is
+                  determined by the number of energy levels as well as the level
+                  information available for radiative decays and collisions.
         """
-        try:
-            return len(self.levels)
-        except MissingDatasetException:
+        # There is no atomic model if there are no energies, collisions, and decays
+        # NOTE: This logic is here rather than using a decorator so that it returns
+        # zero rather than throwing an exception.
+        if not all([self._has_dataset('elvlc'),
+                    self._has_dataset('wgfa'),
+                    self._has_dataset('scups')]):
             return 0
+        n_elvlc = self._elvlc['level'].max()
+        n_wgfa = max(self._wgfa['lower_level'].max(), self._wgfa['upper_level'].max())
+        n_scups = max(self._scups['upper_level'].max(), self._scups['lower_level'].max())
+        # If there is autoionization data associated with this ion, ensure that the model
+        # has enough levels to include these rates.
+        if self._has_dataset('auto'):
+            n_scups = max(n_scups, self._auto['upper_level'].max())
+        return np.min([n_elvlc, n_wgfa, n_scups])
 
     @property
     def n_transitions(self):
@@ -215,7 +234,7 @@ def previous_ion(self):
     @needs_dataset('elvlc', 'wgfa')
     def transitions(self):
         "A `~fiasco.Transitions` object holding the information about transitions for this ion."
-        return Transitions(self.levels, self._wgfa)
+        return Transitions(self.levels, self._wgfa, n_levels=self.n_levels)
 
     @property
     @u.quantity_input
@@ -409,7 +428,7 @@ def effective_collision_strength(self) -> u.dimensionless_unscaled:
         return upsilon.T
 
     @cached_property
-    @needs_dataset('elvlc', 'scups')
+    @needs_dataset('scups')
     @u.quantity_input
     def electron_collision_deexcitation_rate(self) -> u.cm**3 / u.s:
         r"""
@@ -434,11 +453,12 @@ def electron_collision_deexcitation_rate(self) -> u.cm**3 / u.s:
         """
         c = const.h**2 / (2. * np.pi * const.m_e)**(1.5)
         upsilon = self.effective_collision_strength
-        omega_upper = 2. * self._elvlc['J'][self._scups['upper_level'] - 1] + 1.
+        idx = vectorize_where(self.levels.level, self._scups['upper_level'])
+        omega_upper = 2. * self.levels.total_angular_momentum[idx] + 1.
         return c * upsilon / np.sqrt(self.thermal_energy[:, np.newaxis]) / omega_upper
 
     @cached_property
-    @needs_dataset('elvlc', 'scups')
+    @needs_dataset('scups')
     @u.quantity_input
     def electron_collision_excitation_rate(self) -> u.cm**3 / u.s:
         r"""
@@ -463,8 +483,9 @@ def electron_collision_excitation_rate(self) -> u.cm**3 / u.s:
         --------
         electron_collision_deexcitation_rate : De-excitation rate due to collisions
         """
-        omega_upper = 2. * self._elvlc['J'][self._scups['upper_level'] - 1] + 1.
-        omega_lower = 2. * self._elvlc['J'][self._scups['lower_level'] - 1] + 1.
+        J = self.levels.total_angular_momentum
+        omega_upper = 2. * J[vectorize_where(self.levels.level, self._scups['upper_level'])] + 1.
+        omega_lower = 2. * J[vectorize_where(self.levels.level, self._scups['lower_level'])] + 1.
         kBTE = np.outer(1./self.thermal_energy, self._scups['delta_energy'])
         return omega_upper / omega_lower * self.electron_collision_deexcitation_rate * np.exp(-kBTE)
 
@@ -497,7 +518,7 @@ def proton_collision_excitation_rate(self) -> u.cm**3 / u.s:
             return u.Quantity(np.where(ex_rate > 0., ex_rate, 0.), u.cm**3/u.s).T
 
     @cached_property
-    @needs_dataset('elvlc', 'psplups')
+    @needs_dataset('psplups')
     @u.quantity_input
     def proton_collision_deexcitation_rate(self) -> u.cm**3 / u.s:
         r"""
@@ -521,8 +542,9 @@ def proton_collision_deexcitation_rate(self) -> u.cm**3 / u.s:
         proton_collision_excitation_rate : Excitation rate due to collisions with protons
         """
         kBTE = np.outer(self.thermal_energy, 1.0 / self._psplups['delta_energy'])
-        omega_upper = 2. * self._elvlc['J'][self._psplups['upper_level'] - 1] + 1.
-        omega_lower = 2. * self._elvlc['J'][self._psplups['lower_level'] - 1] + 1.
+        J = self.levels.total_angular_momentum
+        omega_upper = 2. * J[vectorize_where(self.levels.level, self._psplups['upper_level'])] + 1.
+        omega_lower = 2. * J[vectorize_where(self.levels.level, self._psplups['lower_level'])] + 1.
         dex_rate = (omega_lower / omega_upper) * self.proton_collision_excitation_rate * np.exp(1. / kBTE)
 
         return dex_rate
@@ -568,7 +590,9 @@ def level_populations(self,
         `~astropy.units.Quantity`
             A ``(l, m, n)`` shaped quantity, where ``l`` is the number of
             temperatures, ``m`` is the number of densities, and ``n`` is the number of energy
-            levels. If ``couple_density_to_temperature=True``, then ``m=1`` and ``l``
+            levels in the ion model. Note that ``n`` will always be the same as the `~fiasco.Ion.n_levels`,
+            but may be different than the number of levels returned by `~fiasco.Levels`.
+            If ``couple_density_to_temperature=True``, then ``m=1`` and ``l``
             represents the number of temperatures and densities.
         """
         if use_two_ion_model:
@@ -695,23 +719,31 @@ def _rate_matrix_radiative_decay(self) -> u.Unit('s-1'):
     @u.quantity_input
     def _rate_matrix_collisional_electron(self) -> u.Unit('cm3 s-1'):
         rate_matrix = u.Quantity(np.zeros(self.temperature.shape + (self.n_levels, self.n_levels,)), 'cm3 s-1')
-        lower_index = self._scups['lower_level'] - 1
-        upper_index = self._scups['upper_level'] - 1
+        # NOTE: For some ions, there may be more rate data available than
+        # there are levels in the model.
+        idx = np.where(np.logical_and(self._scups['lower_level']<=self.n_levels,
+                                      self._scups['upper_level']<=self.n_levels))
+        lower_index = self._scups['lower_level'][idx] - 1
+        upper_index = self._scups['upper_level'][idx] - 1
         # De-excitation from upper states
-        rate_matrix[:, lower_index, upper_index] += self.electron_collision_deexcitation_rate
+        rate_matrix[:, lower_index, upper_index] += self.electron_collision_deexcitation_rate[..., *idx]
         # Excitation from lower states
-        rate_matrix[:, upper_index, lower_index] += self.electron_collision_excitation_rate
+        rate_matrix[:, upper_index, lower_index] += self.electron_collision_excitation_rate[..., *idx]
         return rate_matrix
 
     @cached_property
     @needs_dataset('psplups')
     @u.quantity_input
     def _rate_matrix_collisional_proton(self) -> u.Unit('cm3 s-1'):
         rate_matrix = u.Quantity(np.zeros(self.temperature.shape + (self.n_levels, self.n_levels,)), 'cm3 s-1')
-        lower_index = self._psplups['lower_level'] - 1
-        upper_index = self._psplups['upper_level'] - 1
-        rate_matrix[:, lower_index, upper_index] += self.proton_collision_deexcitation_rate
-        rate_matrix[:, upper_index, lower_index] += self.proton_collision_excitation_rate
+        # NOTE: For some ions, there may be more rate data available than
+        # there are levels in the model.
+        idx = np.where(np.logical_and(self._psplups['lower_level']<=self.n_levels,
+                                      self._psplups['upper_level']<=self.n_levels))
+        lower_index = self._psplups['lower_level'][idx] - 1
+        upper_index = self._psplups['upper_level'][idx] - 1
+        rate_matrix[:, lower_index, upper_index] += self.proton_collision_deexcitation_rate[..., *idx]
+        rate_matrix[:, upper_index, lower_index] += self.proton_collision_excitation_rate[..., *idx]
         return rate_matrix
 
     def _empty_rate_matrix(self, temperature_dependent=True, unit='cm3 s-1'):
@@ -835,10 +867,10 @@ def _rate_matrix_dielectronic_recombination(self) -> u.Unit('cm3 s-1'):
                                              self._auto['upper_level'][idx_ground],
                                              self._auto['autoionization_rate'])
             # Sum radiative decay rates between upper levels and lower bound levels
-            idx_bound = self._wgfa['lower_level'] < self._auto['upper_level'].min()
-            A_rad_sum = vectorize_where_sum(self._wgfa['upper_level'][idx_bound],
+            idx_bound = self.transitions.lower_level < self._auto['upper_level'].min()
+            A_rad_sum = vectorize_where_sum(self.transitions.upper_level[idx_bound],
                                             self._auto['upper_level'][idx_ground],
-                                            self._wgfa['A'][idx_bound],)
+                                            self.transitions.A[idx_bound],)
             branching_ratio = A_rad_sum / (A_rad_sum + A_auto_sum)
             # Get needed levels for recombined and recombining ions
             dc_rate = self._dielectronic_capture_rate(self._auto['lower_level'][idx_ground],
@@ -1112,8 +1144,13 @@ def contribution_function(self, density: u.cm**(-3), **kwargs) -> u.cm**3 * u.er
         upper_level = self.transitions.upper_level[self.transitions.is_bound_bound]
         wavelength = self.transitions.wavelength[self.transitions.is_bound_bound]
         A = self.transitions.A[self.transitions.is_bound_bound]
-        energy = const.h * const.c / wavelength
-        i_upper = vectorize_where(self._elvlc['level'], upper_level)
+        energy = wavelength.to('erg', equivalencies=u.equivalencies.spectral())
+        # NOTE: The first array below provides the correspondence between the last
+        # dimension of the level populations array and the energy level index of
+        # the model. The upper level of the transition is used to make this selection
+        # because the contribution function is proportional to the population of the
+        # level from which the transition is happening.
+        i_upper = vectorize_where(np.arange(1, self.n_levels+1), upper_level)
         g = term * populations[:, :, i_upper] * (A * energy)
         return g
 
@@ -1934,7 +1971,6 @@ def _karzas_cross_section(self, photon_energy: u.erg, ionization_energy: u.erg,
         cross_section[np.where(photon_energy < ionization_energy)] = 0.*cross_section.unit
         return cross_section
 
-    @needs_dataset('elvlc')
     @u.quantity_input
     def two_photon(self,
                    wavelength: u.angstrom,
@@ -1973,22 +2009,16 @@ def two_photon(self,
             A_ji = self._hseq['A']
             psi_norm = self._hseq['psi_norm']
             x_interp, y_interp = self._hseq['y'], self._hseq['psi']
-            config = '2s'  # Get the index of the 2S1/2 state for H-like
-            J = 0.5
+            label = '2s 2S1/2'
         elif self.helium_like:
             A_ji = self._heseq['A']
             psi_norm = 1.0 * u.dimensionless_unscaled
             x_interp, y_interp = self._heseq['y'], self._heseq['psi']
-            config = '1s.2s'  # Get the index of the 1s2s 1S0 state for He-like:
-            J = 0
+            label = '1s 2s 1S0'
         else:
             return u.Quantity(np.zeros(final_shape),  'erg cm^3 s^-1 Angstrom^-1')
-        level_index = np.where((self._elvlc['config'] == config) & (np.isclose(self._elvlc['J'], J)) )[0][0]
-
-        E_obs = self._elvlc['E_obs'][level_index]
-        E_th = self._elvlc['E_th'][level_index]
-        E_2p = E_obs if E_obs > 0.0 else E_th
-        rest_wavelength = 1 / E_2p
+        level_index = np.where(self.levels.label==label)
+        rest_wavelength = self.levels.energy[level_index].to('AA', equivalencies=u.equivalencies.spectral())
 
         # NOTE: Explicitly setting the boundary condition type here to match the behavior of the
         # IDL spline interpolation functions. See https://github.com/wtbarnes/fiasco/pull/297 for
@@ -2004,7 +2034,7 @@ def two_photon(self,
         # for more details.
         kwargs.setdefault('include_protons', False)
         level_population = self.level_populations(electron_density, **kwargs)
-        level_population = level_population[..., level_index]
+        level_population = level_population[..., self.levels[level_index].level-1]
 
         if couple_density_to_temperature:
             electron_density = electron_density[:, np.newaxis]