upgrading env in 2022 (tardis-sn#2410)

* fix several bugs related to upgrade

* fix the get function in pandas

* fixed several incompatibilities for new libraries

* fixed up the changes in pandas testing and other reindexing issues

* fix other bugs with new environment

* blackify tardis

* updated the conda-lock file

* add lock files

* update cache number

* fix simulation tests

* remove units for testing

* blackify

* Update tardis/visualization/tools/tests/test_sdec_plot.py

good catch

Co-authored-by: Atharva Arya <[email protected]>

* add assert

* reset caches

---------

Co-authored-by: Atharva Arya <[email protected]>

tardis/io/atom_data/base.py

@@ -168,7 +168,7 @@ def from_hdf(cls, fname=None):
 
             for name in cls.hdf_names:
                 try:
-                    dataframes[name] = store[name]
+                    dataframes[name] = store.select(name)
                 except KeyError:
                     logger.debug(f"Dataframe does not contain {name} column")
                     nonavailable.append(name)
@@ -287,25 +287,33 @@ def __init__(
         if u.u.cgs == const.u.cgs:
             atom_data.loc[:, "mass"] = Quantity(
                 atom_data["mass"].values, "u"
-            ).cgs
+            ).cgs.value
         else:
-            atom_data.loc[:, "mass"] = atom_data["mass"].values * const.u.cgs
+            atom_data.loc[:, "mass"] = (
+                atom_data["mass"].values * const.u.cgs.value
+            )
 
         # Convert ionization energies to CGS
         ionization_data = ionization_data.squeeze()
-        ionization_data[:] = Quantity(ionization_data[:], "eV").cgs
+        ionization_data[:] = Quantity(ionization_data[:], "eV").cgs.value
 
         # Convert energy to CGS
-        levels.loc[:, "energy"] = Quantity(levels["energy"].values, "eV").cgs
+        levels.loc[:, "energy"] = Quantity(
+            levels["energy"].values, "eV"
+        ).cgs.value
 
         # Create a new columns with wavelengths in the CGS units
-        lines["wavelength_cm"] = Quantity(lines["wavelength"], "angstrom").cgs
+        lines["wavelength_cm"] = Quantity(
+            lines["wavelength"], "angstrom"
+        ).cgs.value
 
         # SET ATTRIBUTES
 
         self.atom_data = atom_data
         self.ionization_data = ionization_data
         self.levels = levels
+        # Not sure why this is need - WEK 17 Sep 2023
+        self.levels.energy = self.levels.energy.astype(np.float64)
         self.lines = lines
 
         # Rename these (drop "_all") when `prepare_atom_data` is removed!
@@ -522,7 +530,9 @@ def prepare_atom_data(
                 self.macro_atom_data.loc[:, "destination_level_idx"] = -1
 
             if self.yg_data is not None:
-                self.yg_data = self.yg_data.loc[self.selected_atomic_numbers]
+                self.yg_data = self.yg_data.reindex(
+                    self.selected_atomic_numbers, level=0
+                )
 
         self.nlte_data = NLTEData(self, nlte_species)
 

tardis/io/model/readers/artis.py

@@ -1,6 +1,6 @@
 import numpy as np
 from astropy import units as u
-from numpy import recfromtxt
+import pandas as pd
 
 
 def read_artis_density(fname):
@@ -43,13 +43,15 @@ def read_artis_density(fname):
         "fe52_fraction",
         "cr48_fraction",
     ]
-    artis_model = recfromtxt(
+
+    artis_model = pd.read_csv(
         fname,
-        skip_header=2,
+        skiprows=2,
         usecols=(0, 1, 2, 4, 5, 6, 7),
-        unpack=True,
-        dtype=[(item, np.float64) for item in artis_model_columns],
-    )
+        dtype={item: np.float64 for item in artis_model_columns},
+        names=artis_model_columns,
+        delim_whitespace=True,
+    ).to_records(index=False)
 
     velocity = u.Quantity(artis_model["velocities"], "km/s").to("cm/s")
     mean_density = u.Quantity(10 ** artis_model["mean_densities_0"], "g/cm^3")[

tardis/io/tests/test_HDFWriter.py

@@ -1,11 +1,8 @@
-import os
-
 import numpy as np
 import pandas as pd
-import pandas.util.testing as pdt
 import pytest
 from astropy import units as u
-from astropy.tests.helper import assert_quantity_allclose
+
 from numpy.testing import assert_almost_equal, assert_array_almost_equal
 
 from tardis.io.util import HDFWriterMixin
@@ -61,6 +58,7 @@ def test_complex_obj_write(tmpdir, attr):
     actual = MockHDF(attr)
     actual.to_hdf(fname, path="test", overwrite=True)
     expected = pd.read_hdf(fname, key="/test/mock_hdf/property").values
+
     assert_array_almost_equal(actual.property, expected)
 
 
@@ -79,7 +77,8 @@ def test_multi_index_write(tmpdir):
     actual.to_hdf(fname, path="test", overwrite=True)
     expected = pd.read_hdf(fname, key="/test/mock_hdf/property")
     expected = pd.MultiIndex.from_tuples(expected.unstack().values)
-    pdt.assert_almost_equal(actual.property, expected)
+    # These are multiindex objects, so we need to compare the values
+    assert np.all(expected.values == actual.property.values)
 
 
 # Test Quantity Objects

tardis/io/tests/test_atomic.py

@@ -49,7 +49,7 @@ def test_atom_data_levels(levels):
 
 def test_atom_data_lines(lines):
     assert_quantity_allclose(
-        lines.at[(2, 0, 0, 6), "wavelength_cm"] * u.Unit("cm"),
+        lines.loc[(2, 0, 0, 6), "wavelength_cm"].values[0] * u.Unit("cm"),
         584.335 * u.Unit("Angstrom"),
     )
 

tardis/model/base.py

@@ -593,7 +593,7 @@ def from_config(cls, config, atom_data=None):
         #       v boundaries.
         no_of_shells = len(velocity) - 1
 
-        if temperature:
+        if temperature is not None:
             t_radiative = temperature
         elif config.plasma.initial_t_rad > 0 * u.K:
             t_radiative = (

tardis/montecarlo/montecarlo_numba/base.py

@@ -1,8 +1,5 @@
 from numba import prange, njit, objmode
-from numba.np.ufunc.parallel import (
-    _get_thread_id as get_thread_id,
-    get_num_threads,
-)
+from numba.np.ufunc.parallel import get_thread_id, get_num_threads
 
 import numpy as np
 

tardis/montecarlo/montecarlo_numba/formal_integral.py

@@ -454,7 +454,7 @@ def make_source_function(self):
         upper_level_index = self.atomic_data.lines.index.droplevel(
             "level_number_lower"
         )
-        e_dot_lu = pd.DataFrame(Edotlu, index=upper_level_index)
+        e_dot_lu = pd.DataFrame(Edotlu.values, index=upper_level_index)
         e_dot_u = e_dot_lu.groupby(level=[0, 1, 2]).sum()
         e_dot_u_src_idx = macro_ref.loc[e_dot_u.index].references_idx.values
 

tardis/montecarlo/tests/test_spectrum.py

@@ -38,7 +38,7 @@ def spectrum(request):
         request.param["lum"],
     )
     distance = request.param["distance"]
-    if distance:
+    if distance is not None:
         data.distance = distance
     return data
 
@@ -67,7 +67,7 @@ def test_luminosity_density_lambda(spectrum):
 
 
 def test_flux_nu(spectrum):
-    if getattr(spectrum, "distance", None):
+    if getattr(spectrum, "distance", None) is not None:
 
         with pytest.warns(DeprecationWarning):
             test_helper.assert_quantity_allclose(
@@ -82,7 +82,7 @@ def test_flux_nu(spectrum):
 
 
 def test_flux_lambda(spectrum):
-    if getattr(spectrum, "distance", None):
+    if getattr(spectrum, "distance", None) is not None:
 
         with pytest.warns(DeprecationWarning):
             test_helper.assert_quantity_allclose(

tardis/plasma/properties/ion_population.py

@@ -53,7 +53,6 @@ class PhiSahaLTE(ProcessingPlasmaProperty):
 
     @staticmethod
     def calculate(g_electron, beta_rad, partition_function, ionization_data):
-
         phis = np.empty(
             (
                 partition_function.shape[0]
@@ -70,9 +69,9 @@ def calculate(g_electron, beta_rad, partition_function, ionization_data):
             current_phis = current_block[1:] / current_block[:-1]
             phis[start_id - i : end_id - i - 1] = current_phis
 
-        broadcast_ionization_energy = ionization_data[
+        broadcast_ionization_energy = ionization_data.reindex(
             partition_function.index
-        ].dropna()
+        ).dropna()
         phi_index = broadcast_ionization_energy.index
         broadcast_ionization_energy = broadcast_ionization_energy.values
 

tardis/plasma/standard_plasmas.py

@@ -278,7 +278,7 @@ def assemble_plasma(config, model, atom_data=None):
     else:
         plasma_modules += helium_lte_properties
 
-    if model._electron_densities:
+    if model._electron_densities is not None:
         electron_densities = pd.Series(model._electron_densities.cgs.value)
         if config.plasma.helium_treatment == "numerical-nlte":
             property_kwargs[IonNumberDensityHeNLTE] = dict(

tardis/plasma/tests/test_complete_plasmas.py

@@ -2,8 +2,10 @@
 import warnings
 
 import pandas as pd
+import pandas.testing as pdt
 import pytest
-from pandas.util import testing as pdt
+
+import numpy as np
 from numpy.testing import assert_almost_equal
 from tardis.io.configuration.config_reader import Configuration
 from tardis.simulation import Simulation
@@ -169,21 +171,29 @@ def plasma(self, request, chianti_he_db_fpath, config, tardis_ref_data):
     def test_plasma_properties(self, plasma, tardis_ref_data, config, attr):
         if hasattr(plasma, attr):
             actual = getattr(plasma, attr)
+            if hasattr(actual, "unit"):
+                actual = actual.value
             if actual.ndim == 1:
                 actual = pd.Series(actual)
             else:
                 actual = pd.DataFrame(actual)
             key = os.path.join(config.plasma.save_path, "plasma", attr)
             expected = tardis_ref_data[key]
-            pdt.assert_almost_equal(actual, expected)
+            if type(actual) == pd.DataFrame:
+                pdt.assert_frame_equal(actual, expected)
+            elif type(actual) == pd.Series:
+                pdt.assert_series_equal(actual, expected)
+            else:
+                raise TypeError(f"Unexpected type {type(actual)}")
+            # we used this before - assert_almost_equal(actual.values, expected.values)
         else:
             warnings.warn(f'Property "{attr}" not found')
 
     def test_levels(self, plasma, tardis_ref_data, config):
         actual = pd.DataFrame(plasma.levels)
         key = os.path.join(config.plasma.save_path, "plasma", "levels")
         expected = tardis_ref_data[key]
-        pdt.assert_almost_equal(actual, expected)
+        pdt.assert_frame_equal(actual, expected)
 
     @pytest.mark.parametrize("attr", scalars_properties)
     def test_scalars_properties(self, plasma, tardis_ref_data, config, attr):
@@ -192,7 +202,7 @@ def test_scalars_properties(self, plasma, tardis_ref_data, config, attr):
             actual = actual.cgs.value
         key = os.path.join(config.plasma.save_path, "plasma", "scalars")
         expected = tardis_ref_data[key][attr]
-        pdt.assert_almost_equal(actual, expected)
+        assert_almost_equal(actual, expected)
 
     def test_helium_treatment(self, plasma, tardis_ref_data, config):
         actual = plasma.helium_treatment

tardis/plasma/tests/test_hdf_plasma.py

@@ -1,8 +1,9 @@
 import os
 import pandas as pd
 import pytest
-import pandas.util.testing as pdt
+
 from numpy.testing import assert_almost_equal
+import pandas.testing as pdt
 from tardis.plasma.properties import property_collections
 
 ###
@@ -68,7 +69,7 @@ def test_hdf_levels(hdf_file_path, simulation_verysimple):
         actual = actual.cgs.value
     path = os.path.join("plasma", "levels")
     expected = pd.read_hdf(hdf_file_path, path)
-    pdt.assert_almost_equal(pd.DataFrame(actual), expected)
+    pdt.assert_frame_equal(pd.DataFrame(actual), expected)
 
 
 scalars_list = ["time_explosion", "link_t_rad_t_electron"]

tardis/simulation/tests/test_simulation.py

@@ -8,10 +8,11 @@
 from tardis.io.configuration.config_reader import Configuration
 from tardis.simulation import Simulation
 from tardis import run_tardis
+import pandas.testing as pdt
 
 import numpy as np
 import pandas as pd
-import pandas.util.testing as pdt
+
 import astropy.units as u
 import tardis
 
@@ -88,10 +89,12 @@ def test_plasma_estimates(simulation_one_loop, refdata, name):
         actual = getattr(simulation_one_loop.transport, name)
     except AttributeError:
         actual = getattr(simulation_one_loop.model, name)
-
+    if name in ["t_radiative", "output_nu", "output_energy"]:
+        # removing the quantitiness of the data
+        actual = actual.value
     actual = pd.Series(actual)
 
-    pdt.assert_almost_equal(actual, refdata(name))
+    pdt.assert_series_equal(actual, refdata(name), rtol=1e-5, atol=1e-8)
 
 
 @pytest.mark.parametrize(
@@ -111,7 +114,12 @@ def test_plasma_state_iterations(simulation_one_loop, refdata, name):
     except Exception:
         actual = pd.DataFrame(actual)
 
-    pdt.assert_almost_equal(actual, refdata(name))
+    if type(actual) == pd.DataFrame:
+        pdt.assert_frame_equal(actual, refdata(name), rtol=1e-5, atol=1e-8)
+    elif type(actual) == pd.Series:
+        pdt.assert_series_equal(actual, refdata(name))
+    else:
+        raise ValueError(f"Unknown type of actual {type(actual)}")
 
 
 @pytest.fixture(scope="module")

tardis/tests/test_tardis_full_formal_integral.py

@@ -90,11 +90,6 @@ def test_spectrum_integrated(self, transport, refdata):
             refdata("spectrum_integrated/luminosity"), "erg /s"
         )
 
-        print(
-            "actual, desired: ",
-            luminosity,
-            transport.spectrum_integrated.luminosity,
-        )
         assert_quantity_allclose(
             transport.spectrum_integrated.luminosity, luminosity
         )

tardis/visualization/tools/sdec_plot.py

@@ -596,7 +596,7 @@ def _calculate_plotting_data(
         self.plot_frequency = self.data[packets_mode].spectrum_frequency
 
         # Filter their plottable range based on packet_wvl_range specified
-        if packet_wvl_range:
+        if packet_wvl_range is not None:
             packet_nu_range = packet_wvl_range.to("Hz", u.spectral())
 
             # Index of value just before the 1st value that is > packet_nu_range[1]
@@ -1224,7 +1224,7 @@ def generate_plot_mpl(
         # Set legends and labels
         self.ax.legend(fontsize=12)
         self.ax.set_xlabel(r"Wavelength $[\mathrm{\AA}]$", fontsize=12)
-        if distance:  # Set y-axis label for flux
+        if distance is not None:  # Set y-axis label for flux
             self.ax.set_ylabel(
                 r"$F_{\lambda}$ [erg $\mathrm{s^{-1}}$ $\mathrm{cm^{-2}}$ $\mathrm{\AA^{-1}}$]",
                 fontsize=12,
@@ -1632,7 +1632,7 @@ def generate_plot_ply(
 
         # Set label and other layout options
         xlabel = pu.axis_label_in_latex("Wavelength", u.AA)
-        if distance:  # Set y-axis label for flux
+        if distance is not None:  # Set y-axis label for flux
             ylabel = pu.axis_label_in_latex(
                 "F_{\\lambda}", u.Unit("erg/(s cm**2 AA)"), only_text=False
             )

tardis/visualization/tools/tests/test_sdec_plot.py

@@ -408,12 +408,14 @@ def test_generate_plot_mpl(
         species_list : list of str
         """
         subgroup_name = make_valid_name("mpl" + request.node.callspec.id)
+        if distance is None:
+            observed_spectrum = None
         fig = plotter.generate_plot_mpl(
             packets_mode=packets_mode,
             packet_wvl_range=packet_wvl_range,
             distance=distance,
             show_modeled_spectrum=show_modeled_spectrum,
-            observed_spectrum=observed_spectrum if distance else None,
+            observed_spectrum=observed_spectrum,
             nelements=nelements,
             species_list=species_list,
         )
@@ -549,12 +551,16 @@ def test_generate_plot_ply(
         species_list : list of str
         """
         subgroup_name = make_valid_name("ply" + request.node.callspec.id)
+        if distance is not None:
+            observed_spectrum = observed_spectrum
+        else:
+            observed_spectrum = None
         fig = plotter.generate_plot_ply(
             packets_mode=packets_mode,
             packet_wvl_range=packet_wvl_range,
             distance=distance,
             show_modeled_spectrum=show_modeled_spectrum,
-            observed_spectrum=observed_spectrum if distance else None,
+            observed_spectrum=observed_spectrum,
             nelements=nelements,
             species_list=species_list,
         )