Switch default get_modeldata() DataFrame to polars and use deprecatio…

…n warnings (#288)

artistools/__init__.py

@@ -34,7 +34,7 @@
 from artistools.inputmodel import get_mean_cell_properties_of_angle_bin as get_mean_cell_properties_of_angle_bin
 from artistools.inputmodel import get_mgi_of_velocity_kms as get_mgi_of_velocity_kms
 from artistools.inputmodel import get_modeldata as get_modeldata
-from artistools.inputmodel import get_modeldata_polars as get_modeldata_polars
+from artistools.inputmodel import get_modeldata_pandas as get_modeldata_pandas
 from artistools.inputmodel import get_modeldata_tuple as get_modeldata_tuple
 from artistools.inputmodel import save_initelemabundances as save_initelemabundances
 from artistools.inputmodel import save_modeldata as save_modeldata

artistools/estimators/estimators.py

@@ -21,6 +21,7 @@
 import pandas as pd
 import polars as pl
 from polars import selectors as cs
+from typing_extensions import deprecated
 
 import artistools as at
 
@@ -403,6 +404,7 @@ def scan_estimators(
     return pldflazy.with_columns(nntot=pl.sum_horizontal(cs.starts_with("nnelement_"))).fill_null(0)
 
 
+@deprecated("Use scan_estimators instead.")
 def read_estimators(
     modelpath: Path | str = Path(),
     modelgridindex: int | Sequence[int] | None = None,

artistools/estimators/plot3destimators_classic.py

@@ -13,7 +13,7 @@
 def read_selected_mgi(
     modelpath: Path, readonly_mgi: list[int] | None = None, readonly_timestep: list[int] | None = None
 ) -> dict[tuple[int, int], t.Any] | None:
-    modeldata, _ = at.inputmodel.get_modeldata(modelpath)
+    modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath)
     return at.estimators.estimators_classic.read_classic_estimators(
         modelpath, modeldata, readonly_mgi=readonly_mgi, readonly_timestep=readonly_timestep
     )

artistools/estimators/plotestimators.py

@@ -58,7 +58,7 @@ def plot_init_abundances(
     assert len(xlist) == len(mgilist)
 
     if seriestype == "initabundances":
-        mergemodelabundata, _ = at.inputmodel.get_modeldata(modelpath, get_elemabundances=True)
+        mergemodelabundata, _ = at.inputmodel.get_modeldata_pandas(modelpath, get_elemabundances=True)
     elif seriestype == "initmasses":
         mergemodelabundata = at.inputmodel.plotinitialcomposition.get_model_abundances_Msun_1D(modelpath)
     else:
@@ -259,7 +259,7 @@ def plot_levelpop(
     else:
         raise ValueError
 
-    modeldata, _ = at.inputmodel.get_modeldata(modelpath, derived_cols=["mass_g", "volume"])
+    modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath, derived_cols=["mass_g", "volume"])
 
     adata = at.atomic.get_levels(modelpath)
 
@@ -1012,7 +1012,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
     if args.classicartis:
         import artistools.estimators.estimators_classic
 
-        modeldata, _ = at.inputmodel.get_modeldata(modelpath)
+        modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath)
         estimatorsdict = artistools.estimators.estimators_classic.read_classic_estimators(modelpath, modeldata)
         assert estimatorsdict is not None
         estimators = pl.DataFrame([
@@ -1070,7 +1070,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
         if not args.x:
             args.x = "velocity"
 
-        dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["ALL"])
+        dfmodel, modelmeta = at.inputmodel.get_modeldata(modelpath, derived_cols=["ALL"])
         if args.x == "velocity" and modelmeta["vmax_cmps"] > 0.3 * 29979245800:
             args.x = "beta"
 
@@ -1087,7 +1087,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
             estimators = estimators.filter(pl.col("modelgridindex").is_in(args.modelgridindex))
 
         if args.classicartis:
-            modeldata, _ = at.inputmodel.get_modeldata(modelpath)
+            modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath)
             allnonemptymgilist = [
                 modelgridindex
                 for modelgridindex in modeldata.index

artistools/gsinetwork/plotqdotabund.py

@@ -426,9 +426,7 @@ def plot_qdot_abund_modelcells(
 
     # arr_z = [at.get_atomic_number(el) for el in arr_el]
 
-    lzdfmodel, modelmeta = at.inputmodel.get_modeldata_polars(
-        modelpath, derived_cols=["mass_g", "rho", "logrho", "volume"]
-    )
+    lzdfmodel, modelmeta = at.inputmodel.get_modeldata(modelpath, derived_cols=["mass_g", "rho", "logrho", "volume"])
     npts_model = modelmeta["npts_model"]
 
     # these factors correct for missing mass due to skipped shells, and volume error due to Cartesian grid map

artistools/inputmodel/__init__.py

@@ -25,7 +25,7 @@
 )
 from artistools.inputmodel.inputmodel_misc import get_mgi_of_velocity_kms as get_mgi_of_velocity_kms
 from artistools.inputmodel.inputmodel_misc import get_modeldata as get_modeldata
-from artistools.inputmodel.inputmodel_misc import get_modeldata_polars as get_modeldata_polars
+from artistools.inputmodel.inputmodel_misc import get_modeldata_pandas as get_modeldata_pandas
 from artistools.inputmodel.inputmodel_misc import get_modeldata_tuple as get_modeldata_tuple
 from artistools.inputmodel.inputmodel_misc import save_empty_abundance_file as save_empty_abundance_file
 from artistools.inputmodel.inputmodel_misc import save_initelemabundances as save_initelemabundances

artistools/inputmodel/describeinputmodel.py

@@ -75,7 +75,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
 
     assert args is not None
 
-    dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(
+    dfmodel, modelmeta = at.inputmodel.get_modeldata(
         args.inputfile,
         get_elemabundances=not args.noabund,
         printwarningsonly=False,

artistools/inputmodel/downscale3dgrid.py

@@ -17,7 +17,7 @@ def make_downscaled_3d_grid(
     """
     modelpath = Path(modelpath)
 
-    dfmodel, modelmeta = at.get_modeldata(modelpath)
+    dfmodel, modelmeta = at.get_modeldata_pandas(modelpath)
     dfelemabund = at.inputmodel.get_initelemabundances_pandas(modelpath)
 
     inputgridsize = modelmeta["ncoordgridx"]

artistools/inputmodel/energyinputfiles.py

@@ -232,7 +232,7 @@ def make_energy_files(rho, Mtot_grams, outputpath: Path | str, modelpath=None, m
 
 def plot_energy_rate(modelpath):
     times_and_rate, _ = at.inputmodel.energyinputfiles.rprocess_const_and_powerlaw()
-    model, _ = at.inputmodel.get_modeldata(modelpath)
+    model, _ = at.inputmodel.get_modeldata_pandas(modelpath)
     Mtot_grams = model["mass_g"].sum()
     plt.plot(
         times_and_rate["times"], np.array(times_and_rate["nuclear_heating_power"]) * Mtot_grams, color="k", zorder=10

artistools/inputmodel/inputmodel_misc.py

@@ -14,6 +14,7 @@
 import pandas as pd
 import polars as pl
 import polars.selectors as cs
+from typing_extensions import deprecated
 
 from artistools.configuration import get_config
 from artistools.misc import firstexisting
@@ -354,7 +355,7 @@ def vectormatch(vec1: list[float], vec2: list[float]) -> bool:
     return dfmodel, modelmeta
 
 
-def get_modeldata_polars(
+def get_modeldata(
     modelpath: Path | str = Path(),
     get_elemabundances: bool = False,
     derived_cols: list[str] | str | None = None,
@@ -460,12 +461,10 @@ def get_modeldata_polars(
     return dfmodel, modelmeta
 
 
+@deprecated("Use get_modeldata() instead.")
 def get_modeldata_tuple(*args: t.Any, **kwargs: t.Any) -> tuple[pd.DataFrame, float, float]:
-    """Get model from model.txt file.
-
-    DEPRECATED: Use get_modeldata() instead.
-    """
-    dfmodel, modelmeta = get_modeldata(*args, **kwargs)
+    """Return tuple of (dfmodel, t_model_init_days, vmax_cmps) from model.txt file."""
+    dfmodel, modelmeta = get_modeldata_pandas(*args, **kwargs)
 
     return dfmodel, modelmeta["t_model_init_days"], modelmeta["vmax_cmps"]
 
@@ -522,15 +521,16 @@ def get_empty_3d_model(
     return dfmodel, modelmeta
 
 
-def get_modeldata(
+@deprecated("Use get_modeldata() instead.")
+def get_modeldata_pandas(
     modelpath: Path | str = Path(),
     get_elemabundances: bool = False,
     derived_cols: list[str] | str | None = None,
     printwarningsonly: bool = False,
     getheadersonly: bool = False,
 ) -> tuple[pd.DataFrame, dict[t.Any, t.Any]]:
-    """Call get_modeldata_polars() and convert to pandas DataFrame."""
-    pldfmodel, modelmeta = get_modeldata_polars(
+    """Like get_modeldata() but convert polars DataFrame to pandas."""
+    pldfmodel, modelmeta = get_modeldata(
         modelpath=modelpath,
         get_elemabundances=get_elemabundances,
         derived_cols=derived_cols,
@@ -540,7 +540,7 @@ def get_modeldata(
     dfmodel = pldfmodel.collect().to_pandas(use_pyarrow_extension_array=True)
     if modelmeta["npts_model"] > 100000 and not getheadersonly:
         # dfmodel.info(verbose=False, memory_usage="deep")
-        print("WARNING: Using pandas DataFrame for large model data. Switch to using get_modeldata_polars() instead.")
+        print("WARNING: Using pandas DataFrame for large model data. Switch to using get_modeldata() instead.")
 
     return dfmodel, modelmeta
 

artistools/inputmodel/makeartismodel.py

@@ -78,7 +78,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
         ndim_out = args.dimensionreduce
         assert ndim_out in {0, 1, 2}
         for modelpath in args.modelpath:
-            dfmodel, modelmeta = at.inputmodel.get_modeldata_polars(modelpath, derived_cols=["mass_g"])
+            dfmodel, modelmeta = at.inputmodel.get_modeldata(modelpath, derived_cols=["mass_g"])
             ndim_in = modelmeta["dimensions"]
             if ndim_in <= ndim_out:
                 msg = f"Cannot reduce {ndim_in}D model to {ndim_out}D"
@@ -117,7 +117,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
         )
 
     if args.makeenergyinputfiles:
-        model, modelmeta = at.inputmodel.get_modeldata(args.modelpath[0], derived_cols=["mass_g"])
+        model, modelmeta = at.inputmodel.get_modeldata_pandas(args.modelpath[0], derived_cols=["mass_g"])
         rho = 10 ** model["logrho"] if modelmeta["dimensions"] == 1 else model["rho"]
         Mtot_grams = model["mass_g"].sum()
 

artistools/inputmodel/plotdensity.py

@@ -47,9 +47,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None =
         args.modelpath = ["."]
 
     for modelpath in args.modelpath:
-        dfmodel, modelmeta = at.get_modeldata_polars(
-            modelpath, derived_cols=["vel_r_min", "vel_r_mid", "vel_r_max", "mass_g"]
-        )
+        dfmodel, modelmeta = at.get_modeldata(modelpath, derived_cols=["vel_r_min", "vel_r_mid", "vel_r_max", "mass_g"])
         label = at.get_model_name(modelpath)
         print(f"Plotting {label}")
         binned_xvals: list[float] = []

artistools/inputmodel/plotinitialcomposition.py

@@ -122,7 +122,7 @@ def plot_slice_modelcolumn(
 def plot_2d_initial_abundances(modelpath, args: argparse.Namespace) -> None:
     # if the species ends in a number then we need to also get the nuclear mass fractions (not just element abundances)
     get_elemabundances = any(plotvar[-1] not in string.digits for plotvar in args.plotvars)
-    dfmodel, modelmeta = at.get_modeldata(
+    dfmodel, modelmeta = at.get_modeldata_pandas(
         modelpath, get_elemabundances=get_elemabundances, derived_cols=["pos_min", "pos_max"]
     )
     assert modelmeta["dimensions"] > 1
@@ -241,7 +241,7 @@ def get_model_abundances_Msun_1D(modelpath: Path) -> pd.DataFrame:
 
 
 def plot_most_abundant(modelpath, args: argparse.Namespace):  # noqa: ARG001
-    model, _ = at.inputmodel.get_modeldata(modelpath[0])
+    model, _ = at.inputmodel.get_modeldata_pandas(modelpath[0])
     abundances = at.inputmodel.get_initelemabundances_pandas(modelpath[0])
 
     merge_dfs = model.merge(abundances, how="inner", on="inputcellid")
@@ -364,7 +364,7 @@ def make_3d_plot(modelpath, args: argparse.Namespace) -> None:
 
 
 def plot_phi_hist(modelpath):
-    dfmodel, _ = at.get_modeldata(modelpath, derived_cols=["pos_x_mid", "pos_y_mid", "pos_z_mid", "vel_r_mid"])
+    dfmodel, _ = at.get_modeldata_pandas(modelpath, derived_cols=["pos_x_mid", "pos_y_mid", "pos_z_mid", "vel_r_mid"])
     # print(dfmodel.keys())
     # quit()
     at.inputmodel.inputmodel_misc.get_cell_angle(dfmodel, modelpath)

artistools/inputmodel/slice1dfromconein3dmodel.py

@@ -22,7 +22,7 @@ def make_cone(args):
 
     theta = np.radians([angle_of_cone / 2])  # angle between line of sight and edge is half angle of cone
 
-    dfmodel, modelmeta = at.get_modeldata(modelpath=args.modelpath[0], get_elemabundances=True)
+    dfmodel, modelmeta = at.get_modeldata_pandas(modelpath=args.modelpath[0], get_elemabundances=True)
     args.t_model = modelmeta["t_model_init_days"]
 
     if args.positive_axis:
@@ -54,7 +54,9 @@ def get_profile_along_axis(args: argparse.Namespace, modeldata=None, derived_col
 
     # merge_dfs, args.t_model, args.vmax = at.inputmodel.get_modeldata_tuple(args.modelpath, dimensions=3, get_elemabundances=True)
     if modeldata is None:
-        modeldata, _ = at.inputmodel.get_modeldata(args.modelpath, get_elemabundances=True, derived_cols=derived_cols)
+        modeldata, _ = at.inputmodel.get_modeldata_pandas(
+            args.modelpath, get_elemabundances=True, derived_cols=derived_cols
+        )
 
     position_closest_to_axis = modeldata.iloc[(modeldata[f"pos_{args.other_axis2}_min"]).abs().argsort()][:1][
         f"pos_{args.other_axis2}_min"

artistools/inputmodel/test_inputmodel.py

@@ -28,44 +28,46 @@ def test_describeinputmodel_3d() -> None:
 
 def test_get_modeldata_1d() -> None:
     for getheadersonly in (False, True):
-        dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath, getheadersonly=getheadersonly)
+        _, modelmeta = at.get_modeldata(modelpath=modelpath, getheadersonly=getheadersonly)
         assert np.isclose(modelmeta["vmax_cmps"], 800000000.0)
         assert modelmeta["dimensions"] == 1
         assert modelmeta["npts_model"] == 1
 
-    dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath, derived_cols=["mass_g"])
-    assert np.isclose(dfmodel.mass_g.sum(), 1.416963e33)
+    lzdfmodel, modelmeta = at.get_modeldata(modelpath=modelpath, derived_cols=["mass_g"])
+    assert np.isclose(lzdfmodel.select(pl.col("mass_g").sum()).collect().item(), 1.416963e33)
 
 
 @pytest.mark.benchmark
 def test_get_modeldata_3d() -> None:
     for getheadersonly in (False, True):
-        dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath_3d, getheadersonly=getheadersonly)
+        _, modelmeta = at.get_modeldata(modelpath=modelpath_3d, getheadersonly=getheadersonly)
         assert np.isclose(modelmeta["vmax_cmps"], 2892020000.0)
         assert modelmeta["dimensions"] == 3
         assert modelmeta["npts_model"] == 1000
         assert modelmeta["ncoordgridx"] == 10
 
-    dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath_3d, derived_cols=["mass_g"])
-    assert np.isclose(dfmodel.mass_g.sum(), 2.7861855e33)
+    lzdfmodel, modelmeta = at.get_modeldata(modelpath=modelpath_3d, derived_cols=["mass_g"])
+    assert np.isclose(lzdfmodel.select(pl.col("mass_g").sum()).collect().item(), 2.7861855e33)
 
 
 def test_get_cell_angle() -> None:
-    modeldata, _ = at.inputmodel.get_modeldata(
+    modeldata, _ = at.inputmodel.get_modeldata_pandas(
         modelpath=modelpath_3d, derived_cols=["pos_x_mid", "pos_y_mid", "pos_z_mid"]
     )
     at.inputmodel.inputmodel_misc.get_cell_angle(modeldata, modelpath=modelpath_3d)
     assert "cos_bin" in modeldata
 
 
 def test_downscale_3dmodel() -> None:
-    dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath_3d, get_elemabundances=True, derived_cols=["mass_g"])
+    lzdfmodel, modelmeta = at.get_modeldata(modelpath=modelpath_3d, get_elemabundances=True, derived_cols=["mass_g"])
     modelpath_3d_small = at.inputmodel.downscale3dgrid.make_downscaled_3d_grid(
         modelpath_3d, outputgridsize=2, outputfolder=outputpath
     )
-    dfmodel_small, modelmeta_small = at.get_modeldata(
+    dfmodel = lzdfmodel.collect()
+    lzdfmodel_small, modelmeta_small = at.get_modeldata(
         modelpath_3d_small, get_elemabundances=True, derived_cols=["mass_g"]
     )
+    dfmodel_small = lzdfmodel_small.collect()
     assert np.isclose(dfmodel["mass_g"].sum(), dfmodel_small["mass_g"].sum())
     assert np.isclose(modelmeta["vmax_cmps"], modelmeta_small["vmax_cmps"])
     assert np.isclose(modelmeta["t_model_init_days"], modelmeta_small["t_model_init_days"])
@@ -178,11 +180,11 @@ def test_makeartismodelfrom_sph_particles() -> None:
                     atol=1e-4,
                 )
             else:
-                dfmodel3lz, _ = at.inputmodel.get_modeldata_polars(
+                dfmodel3lz, _ = at.inputmodel.get_modeldata(
                     modelpath=outputpath / f"kilonova_{3:d}d", derived_cols=["mass_g"]
                 )
                 dfmodel3 = dfmodel3lz.collect()
-                dfmodel_lowerdlz, _ = at.inputmodel.get_modeldata_polars(
+                dfmodel_lowerdlz, _ = at.inputmodel.get_modeldata(
                     modelpath=outputpath / f"kilonova_{dimensions:d}d", derived_cols=["mass_g"]
                 )
                 dfmodel_lowerd = dfmodel_lowerdlz.collect()
@@ -323,7 +325,7 @@ def test_save_load_3d_model() -> None:
 
     # first load will be from text, second from parquet
     for _ in (0, 1):
-        dfmodel_loaded, modelmeta_loaded = at.inputmodel.get_modeldata_polars(modelpath=outpath)
+        dfmodel_loaded, modelmeta_loaded = at.inputmodel.get_modeldata(modelpath=outpath)
         pltest.assert_frame_equal(
             dfmodel, dfmodel_loaded.collect(), check_column_order=False, check_dtypes=False, rtol=1e-4, atol=1e-4
         )
@@ -372,9 +374,7 @@ def run_dimension_reduce():
         )
         at.inputmodel.save_modeldata(outpath=outpath, dfmodel=dfmodel_lowerd, modelmeta=modelmeta_lowerd)
 
-    dfmodel_lowerd_lz, _ = at.inputmodel.get_modeldata_polars(
-        modelpath=outpath, derived_cols=["mass_g", "kinetic_en_erg"]
-    )
+    dfmodel_lowerd_lz, _ = at.inputmodel.get_modeldata(modelpath=outpath, derived_cols=["mass_g", "kinetic_en_erg"])
     dfmodel_lowerd = dfmodel_lowerd_lz.collect()
 
     # check that the total mass is conserved

artistools/lightcurve/lightcurve.py

@@ -70,7 +70,7 @@ def get_from_packets(
     arr_timedelta = at.get_timestep_times(modelpath=modelpath, loc="delta")
     # timearray = np.arange(250, 350, 0.1)
     if get_cmf_column:
-        _, modelmeta = at.inputmodel.get_modeldata(modelpath, getheadersonly=True, printwarningsonly=True)
+        _, modelmeta = at.inputmodel.get_modeldata_pandas(modelpath, getheadersonly=True, printwarningsonly=True)
         escapesurfacegamma = math.sqrt(1 - (modelmeta["vmax_cmps"] / 29979245800) ** 2)
     else:
         escapesurfacegamma = None

artistools/lightcurve/plotlightcurve.py

@@ -41,9 +41,7 @@ def plot_deposition_thermalisation(
     #     axistherm.set_ylim(bottom=0.1, top=1.0)
 
     if args.plotthermalisation:
-        dfmodel, _ = at.inputmodel.get_modeldata_polars(
-            modelpath, derived_cols=["mass_g", "vel_r_mid", "kinetic_en_erg"]
-        )
+        dfmodel, _ = at.inputmodel.get_modeldata(modelpath, derived_cols=["mass_g", "vel_r_mid", "kinetic_en_erg"])
 
         model_mass_grams = dfmodel.select("mass_g").sum().collect().item()
         print(f"  model mass: {model_mass_grams / 1.989e33:.3f} Msun")

artistools/linefluxes.py

@@ -148,7 +148,7 @@ def get_line_fluxes_from_pops(emfeatures, modelpath, arr_tstart=None, arr_tend=N
     # arr_timedelta = np.array(arr_tend) - np.array(arr_tstart)
     arr_tmid = arr_tend = (np.array(arr_tstart) + np.array(arr_tend)) / 2.0
 
-    modeldata, _ = at.inputmodel.get_modeldata(modelpath)
+    modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath)
 
     ionlist = [(feature.atomic_number, feature.ion_stage) for feature in emfeatures]
     adata = at.atomic.get_levels(modelpath, ionlist=tuple(ionlist), get_transitions=True, get_photoionisations=False)
@@ -656,7 +656,7 @@ def make_emitting_regions_plot(args: argparse.Namespace) -> None:
                         }
 
             estimators = at.estimators.read_estimators(modelpath)
-            modeldata, _ = at.inputmodel.get_modeldata(modelpath)
+            modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath)
             Tedata_all[modelindex] = {}
             log10nnedata_all[modelindex] = {}
             for tmid, tstart, tend in zip(times_days, args.timebins_tstart, args.timebins_tend, strict=False):