Add functionality to allow user to modify abundances (#258)

* add functionality to allow user to modify abundances

* precommit fixes

* fix typo

* fix abundance test in test_ion.py

* precommit fix, attempt 2

* update setters and fix test_ion

* small fix to test_ion.py

* simplify element setter

* fix _instance_kwargs and add test

* add photospheric abundances to test data

* simplify iteration over ions when setting abundance

* ensure return type is a unitless scalar; set dset to none if abundance is a float

* clean up use of abundance_filename key

* simplify tests

---------

Co-authored-by: Will Barnes <will.t.barnes@gmail.com>

fiasco/base.py

@@ -143,7 +143,7 @@ class IonBase(Base):
     """
 
     @property
-    def _abundance(self):
+    def _abund(self):
         data_path = '/'.join([self.atomic_symbol.lower(), 'abundance'])
         return DataIndexer.create_indexer(self.hdf5_dbase_root, data_path)
 

fiasco/conftest.py

@@ -20,6 +20,7 @@
 TEST_FILES = {
     'sun_coronal_1992_feldman_ext.abund',
     'sun_coronal_1992_feldman.abund',
+    'sun_photospheric_2007_grevesse.abund',
     'chianti.ip',
     'chianti.ioneq',
     'gffgu.dat',

fiasco/elements.py

@@ -61,6 +61,11 @@ def element_name(self):
     def abundance(self):
         return self[0].abundance
 
+    @abundance.setter
+    def abundance(self, abundance):
+        for _ion in self:
+            _ion.abundance = abundance
+
     @cached_property
     def _rate_matrix(self):
         rate_matrix = np.zeros(self.temperature.shape+(self.atomic_number+1, self.atomic_number+1))

fiasco/fiasco.py

@@ -83,7 +83,7 @@ def proton_electron_ratio(temperature: u.K, **kwargs):
         try:
             abundance = el.abundance
         except KeyError:
-            abund_file = el[0]._instance_kwargs['abundance_filename']
+            abund_file = el[0]._instance_kwargs['abundance']
             log.warning(
                 f'Not including {el.atomic_symbol}. Abundance not available from {abund_file}.')
             continue

fiasco/ions.py

@@ -42,23 +42,24 @@ class Ion(IonBase, ContinuumBase):
         Temperature array over which to evaluate temperature dependent quantities.
     ioneq_filename : `str`, optional
         Ionization equilibrium dataset
-    abundance_filename : `str`, optional
-        Abundance dataset
+    abundance : `str` or `float`, optional
+        If a string is provided, use the appropriate abundance dataset.
+        If a float is provided, use that value as the abundance.
     ip_filename : `str`, optional
         Ionization potential dataset
     """
 
     @u.quantity_input
-    def __init__(self, ion_name, temperature: u.K, *args, **kwargs):
+    def __init__(self, ion_name, temperature: u.K,
+                abundance = 'sun_coronal_1992_feldman_ext', *args, **kwargs):
         super().__init__(ion_name, *args, **kwargs)
         self.temperature = np.atleast_1d(temperature)
         # Get selected datasets
         # TODO: do not hardcode defaults, pull from rc file
         self._dset_names = {}
         self._dset_names['ioneq_filename'] = kwargs.get('ioneq_filename', 'chianti')
-        self._dset_names['abundance_filename'] = kwargs.get('abundance_filename',
-                                                            'sun_coronal_1992_feldman_ext')
         self._dset_names['ip_filename'] = kwargs.get('ip_filename', 'chianti')
+        self.abundance = abundance
 
     def _new_instance(self, temperature=None, **kwargs):
         """
@@ -94,7 +95,7 @@ def __repr__(self):
 HDF5 Database: {self.hdf5_dbase_root}
 Using Datasets:
   ioneq: {self._dset_names['ioneq_filename']}
-  abundance: {self._dset_names['abundance_filename']}
+  abundance: {self._dset_names.get('abundance', self.abundance)}
   ip: {self._dset_names['ip_filename']}"""
 
     @cached_property
@@ -134,6 +135,12 @@ def _instance_kwargs(self):
             'hdf5_dbase_root': self.hdf5_dbase_root,
             **self._dset_names,
         }
+        # If the abundance is set using a string specifying the abundance dataset,
+        # the dataset name is in _dset_names. We want to pass this to the new instance
+        # so that the new instance knows that the abundance was specified using a
+        # dataset name. Otherwise, we can just pass the actual abundance value.
+        if kwargs['abundance'] is None:
+            kwargs['abundance'] = self.abundance
         return kwargs
 
     def next_ion(self):
@@ -204,11 +211,26 @@ def ioneq(self):
         return u.Quantity(ioneq)
 
     @property
-    def abundance(self):
+    @u.quantity_input
+    def abundance(self) -> u.dimensionless_unscaled:
         """
         Elemental abundance relative to H.
         """
-        return self._abundance[self._dset_names['abundance_filename']]
+        return self._abundance
+
+    @abundance.setter
+    def abundance(self, abundance):
+        """
+        Sets the abundance of an ion (relative to H).
+        If the abundance is given as a string, use the matching abundance set.
+        If the abundance is given as a float, use that value directly.
+        """
+        if isinstance(abundance, str):
+            self._dset_names['abundance'] = abundance
+            self._abundance = self._abund[self._dset_names['abundance']]
+        else:
+            self._dset_names['abundance'] = None
+            self._abundance = abundance
 
     @property
     @needs_dataset('ip')

fiasco/tests/idl/test_idl_continuum.py

@@ -17,7 +17,7 @@ def all_ions(hdf5_dbase_root):
     abundance_name = 'sun_coronal_1992_feldman_ext'
     ioneq_name = 'chianti'
     ion_kwargs = {
-        'abundance_filename': abundance_name,
+        'abundance': abundance_name,
         'ioneq_filename': ioneq_name,
         'hdf5_dbase_root': hdf5_dbase_root,
     }
@@ -35,7 +35,7 @@ def idl_input_args(all_ions, wavelength):
     return {
         'wavelength': wavelength,
         'temperature': all_ions.temperature,
-        'abundance': all_ions[0]._dset_names['abundance_filename'],
+        'abundance': all_ions[0]._dset_names['abundance'],
         'ioneq': all_ions[0]._dset_names['ioneq_filename'],
     }
 

fiasco/tests/idl/test_idl_goft.py

@@ -72,7 +72,7 @@ def test_idl_compare_goft(idl_env, hdf5_dbase_root, dbase_version, ion_name, wav
     ion = fiasco.Ion(ion_name,
                      idl_result['temperature'],
                      hdf5_dbase_root=hdf5_dbase_root,
-                     abundance_filename=idl_result['abundance'],
+                     abundance=idl_result['abundance'],
                      ioneq_filename=idl_result['ioneq'])
     contribution_func = ion.contribution_function(idl_result['density'])
     idx = np.argmin(np.abs(ion.transitions.wavelength[~ion.transitions.is_twophoton] - idl_result['wavelength']))

fiasco/tests/test_element.py

@@ -83,3 +83,14 @@ def test_equilibrium_ionization(hdf5_dbase_root):
 
 def test_element_repr(element):
     assert element.element_name in element.__repr__()
+
+
+@pytest.mark.parametrize(('value', 'dset'),[
+    (0.07943282347242822, 'sun_coronal_1992_feldman_ext'),
+    (0.08511380382023759, 'sun_photospheric_2007_grevesse'),
+    (1e-3, None),
+])
+def test_change_element_abundance(another_element, value, dset):
+    another_element.abundance = value if dset is None else dset
+    assert u.allclose(another_element.abundance, value)
+    assert u.allclose(another_element[1].abundance, value)

fiasco/tests/test_ion.py

@@ -45,16 +45,16 @@ def fe20(hdf5_dbase_root):
 
 
 def test_new_instance(ion):
-    abundance_filename = ion._instance_kwargs['abundance_filename']
+    abundance = ion._instance_kwargs['abundance']
     new_ion = ion._new_instance()
     for k in new_ion._instance_kwargs:
         assert new_ion._instance_kwargs[k] == ion._instance_kwargs[k]
     assert u.allclose(new_ion.temperature, ion.temperature, rtol=0)
     new_ion = ion._new_instance(temperature=ion.temperature[:1])
     assert u.allclose(new_ion.temperature, ion.temperature[:1])
-    new_ion = ion._new_instance(abundance_filename='sun_coronal_1992_feldman')
-    assert new_ion._instance_kwargs['abundance_filename'] == 'sun_coronal_1992_feldman'
-    assert ion._instance_kwargs['abundance_filename'] == abundance_filename
+    new_ion = ion._new_instance(abundance='sun_coronal_1992_feldman')
+    assert new_ion._instance_kwargs['abundance'] == 'sun_coronal_1992_feldman'
+    assert ion._instance_kwargs['abundance'] == abundance
 
 
 def test_level_indexing(ion):
@@ -137,7 +137,7 @@ def test_ioneq_out_bounds_is_nan(ion):
     assert np.isnan(ion_out_of_bounds.ioneq).all()
 
 
-def test_formation_temeprature(ion):
+def test_formation_temperature(ion):
     assert ion.formation_temperature == ion.temperature[np.argmax(ion.ioneq)]
 
 
@@ -156,13 +156,11 @@ def test_proton_collision(fe10):
 
 
 def test_missing_abundance(hdf5_dbase_root):
-    ion = fiasco.Ion('Li 1',
-                     temperature,
-                     abundance_filename='sun_coronal_1992_feldman',
-                     hdf5_dbase_root=hdf5_dbase_root)
     with pytest.raises(KeyError):
-        _ = ion.abundance
-
+        fiasco.Ion('Li 1',
+                    temperature,
+                    abundance='sun_coronal_1992_feldman',
+                    hdf5_dbase_root=hdf5_dbase_root)
 
 def test_ip(ion):
     assert ion.ip.dtype == np.dtype('float64')
@@ -389,3 +387,29 @@ 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
+
+
+@pytest.mark.parametrize(('value', 'dset'),[
+    (0.0001258925411794166, 'sun_coronal_1992_feldman_ext'),
+    (2.818382931264455e-05, 'sun_photospheric_2007_grevesse'),
+    (1e-3, None),
+])
+def test_change_ion_abundance(ion, value, dset):
+    ion.abundance = value if dset is None else dset
+    assert u.allclose(ion.abundance, value)
+    assert ion._dset_names['abundance'] == dset
+    assert ion._instance_kwargs['abundance'] == (value if dset is None else dset)
+
+
+def test_new_instance_abundance_preserved_float(ion):
+    ion.abundance = 1e-3
+    new_ion = ion._new_instance()
+    assert u.allclose(new_ion.abundance, ion.abundance)
+    assert new_ion._dset_names['abundance'] is None
+
+
+def test_new_instance_abundance_preserved_string(ion):
+    ion.abundance = 'sun_photospheric_2007_grevesse'
+    new_ion = ion._new_instance()
+    assert u.allclose(new_ion.abundance, 2.818382931264455e-05)
+    assert new_ion._dset_names['abundance'] == 'sun_photospheric_2007_grevesse'