From 403eb088891f2c32fac5f7d85e235c020b736afd Mon Sep 17 00:00:00 2001 From: Luke Shingles Date: Mon, 6 Jan 2025 12:26:48 +0000 Subject: [PATCH] Switch default get_modeldata() DataFrame to polars and use deprecation warnings (#288) --- artistools/__init__.py | 2 +- artistools/estimators/estimators.py | 2 ++ .../estimators/plot3destimators_classic.py | 2 +- artistools/estimators/plotestimators.py | 10 +++---- artistools/gsinetwork/plotqdotabund.py | 4 +-- artistools/inputmodel/__init__.py | 2 +- artistools/inputmodel/describeinputmodel.py | 2 +- artistools/inputmodel/downscale3dgrid.py | 2 +- artistools/inputmodel/energyinputfiles.py | 2 +- artistools/inputmodel/inputmodel_misc.py | 20 ++++++------- artistools/inputmodel/makeartismodel.py | 4 +-- artistools/inputmodel/plotdensity.py | 4 +-- .../inputmodel/plotinitialcomposition.py | 6 ++-- .../inputmodel/slice1dfromconein3dmodel.py | 6 ++-- artistools/inputmodel/test_inputmodel.py | 30 +++++++++---------- artistools/lightcurve/lightcurve.py | 2 +- artistools/lightcurve/plotlightcurve.py | 4 +-- artistools/linefluxes.py | 4 +-- artistools/misc.py | 8 ++--- artistools/nltepops/plotnltepops.py | 4 +-- artistools/nonthermal/solvespencerfanocmd.py | 2 +- artistools/packets/packetsplots.py | 2 +- artistools/plotspherical.py | 2 +- artistools/spectra/plotspectra.py | 6 ++-- artistools/spectra/spectra.py | 4 +-- artistools/transitions.py | 2 +- artistools/viewing_angles_visualization.py | 2 +- artistools/writecomparisondata.py | 8 ++--- pyproject.toml | 1 + 29 files changed, 74 insertions(+), 75 deletions(-) diff --git a/artistools/__init__.py b/artistools/__init__.py index 29c2fed01..e24c30f20 100644 --- a/artistools/__init__.py +++ b/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 diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py index 3e6b3eb33..f63ac46e1 100755 --- a/artistools/estimators/estimators.py +++ b/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, diff --git a/artistools/estimators/plot3destimators_classic.py b/artistools/estimators/plot3destimators_classic.py index 200b3e423..cad80140d 100644 --- a/artistools/estimators/plot3destimators_classic.py +++ b/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 ) diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py index e9d7e32fe..e81b147f4 100755 --- a/artistools/estimators/plotestimators.py +++ b/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 diff --git a/artistools/gsinetwork/plotqdotabund.py b/artistools/gsinetwork/plotqdotabund.py index 07443c308..cee5c2da8 100644 --- a/artistools/gsinetwork/plotqdotabund.py +++ b/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 diff --git a/artistools/inputmodel/__init__.py b/artistools/inputmodel/__init__.py index 2343e3415..8a003be89 100644 --- a/artistools/inputmodel/__init__.py +++ b/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 diff --git a/artistools/inputmodel/describeinputmodel.py b/artistools/inputmodel/describeinputmodel.py index d4f1defbf..b551e891d 100755 --- a/artistools/inputmodel/describeinputmodel.py +++ b/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, diff --git a/artistools/inputmodel/downscale3dgrid.py b/artistools/inputmodel/downscale3dgrid.py index 670d87e35..ff5de2a6c 100644 --- a/artistools/inputmodel/downscale3dgrid.py +++ b/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"] diff --git a/artistools/inputmodel/energyinputfiles.py b/artistools/inputmodel/energyinputfiles.py index c36d97133..4722070b9 100644 --- a/artistools/inputmodel/energyinputfiles.py +++ b/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 diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py index 835ba49c2..dc5c717a7 100644 --- a/artistools/inputmodel/inputmodel_misc.py +++ b/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 diff --git a/artistools/inputmodel/makeartismodel.py b/artistools/inputmodel/makeartismodel.py index fef8c5485..7d619cfa2 100755 --- a/artistools/inputmodel/makeartismodel.py +++ b/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() diff --git a/artistools/inputmodel/plotdensity.py b/artistools/inputmodel/plotdensity.py index c8c4bee9a..62a61c914 100755 --- a/artistools/inputmodel/plotdensity.py +++ b/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] = [] diff --git a/artistools/inputmodel/plotinitialcomposition.py b/artistools/inputmodel/plotinitialcomposition.py index 93ca95534..70ca5ae6e 100755 --- a/artistools/inputmodel/plotinitialcomposition.py +++ b/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) diff --git a/artistools/inputmodel/slice1dfromconein3dmodel.py b/artistools/inputmodel/slice1dfromconein3dmodel.py index e6316f864..8e7d54e73 100755 --- a/artistools/inputmodel/slice1dfromconein3dmodel.py +++ b/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" diff --git a/artistools/inputmodel/test_inputmodel.py b/artistools/inputmodel/test_inputmodel.py index 6b6f6ed62..dddca6b73 100644 --- a/artistools/inputmodel/test_inputmodel.py +++ b/artistools/inputmodel/test_inputmodel.py @@ -28,30 +28,30 @@ 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) @@ -59,13 +59,15 @@ def test_get_cell_angle() -> None: 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 diff --git a/artistools/lightcurve/lightcurve.py b/artistools/lightcurve/lightcurve.py index f454cb568..4e58ab75e 100644 --- a/artistools/lightcurve/lightcurve.py +++ b/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 diff --git a/artistools/lightcurve/plotlightcurve.py b/artistools/lightcurve/plotlightcurve.py index 391d558b7..4991ac106 100644 --- a/artistools/lightcurve/plotlightcurve.py +++ b/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") diff --git a/artistools/linefluxes.py b/artistools/linefluxes.py index 5cdddcea8..24e758024 100755 --- a/artistools/linefluxes.py +++ b/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): diff --git a/artistools/misc.py b/artistools/misc.py index 7ea07442c..1a992f25b 100644 --- a/artistools/misc.py +++ b/artistools/misc.py @@ -291,9 +291,9 @@ def get_wid_init_at_tmodel( if ngridpoints is None or t_model_days is None or xmax is None: # Luke: ngridpoint only equals the number of model cells if the model is 3D assert modelpath is not None - from artistools.inputmodel import get_modeldata + from artistools.inputmodel import get_modeldata_pandas - _, modelmeta = get_modeldata(modelpath, getheadersonly=True) + _, modelmeta = get_modeldata_pandas(modelpath, getheadersonly=True) assert modelmeta["dimensions"] == 3 ngridpoints = modelmeta["npts_model"] xmax = modelmeta["vmax_cmps"] * modelmeta["t_model_init_days"] * 86400.0 @@ -549,9 +549,9 @@ def get_timestep_time(modelpath: Path | str, timestep: int) -> float: def get_escaped_arrivalrange(modelpath: Path | str) -> tuple[int, float | None, float | None]: """Return the time range for which the entire model can send light signals the observer.""" modelpath = Path(modelpath) - from artistools.inputmodel import get_modeldata + from artistools.inputmodel import get_modeldata_pandas - _, modelmeta = get_modeldata(modelpath, printwarningsonly=True, getheadersonly=True) + _, modelmeta = get_modeldata_pandas(modelpath, printwarningsonly=True, getheadersonly=True) vmax = modelmeta["vmax_cmps"] cornervmax = math.sqrt(3 * vmax**2) diff --git a/artistools/nltepops/plotnltepops.py b/artistools/nltepops/plotnltepops.py index 405aa1488..2cf6a77b8 100755 --- a/artistools/nltepops/plotnltepops.py +++ b/artistools/nltepops/plotnltepops.py @@ -147,7 +147,7 @@ def get_floers_data(dfpopthision, atomic_number, ion_stage, modelpath, T_e, mode floersmultizonefilename = "level_pops_subch_shen2018-247d.csv" if floersmultizonefilename and Path(floersmultizonefilename).is_file(): - modeldata, _ = at.inputmodel.get_modeldata(modelpath) # TODO: move into modelpath loop + modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath) # TODO: move into modelpath loop vel_outer = modeldata.iloc[modelgridindex].vel_r_max_kmps print(f" reading {floersmultizonefilename}", vel_outer, T_e) dffloers = pd.read_csv(floersmultizonefilename) @@ -470,7 +470,7 @@ def plot_populations_with_time_or_velocity( modelgridindex_list = [int(args.modelgridindex[0])] * len(timesteps) if args.x == "velocity": - modeldata, _ = at.inputmodel.get_modeldata(modelpaths[0]) # TODO: move into modelpath loop + modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpaths[0]) # TODO: move into modelpath loop velocity = modeldata["vel_r_max_kmps"] modelgridindex_list = [mgi for mgi, _ in enumerate(velocity)] diff --git a/artistools/nonthermal/solvespencerfanocmd.py b/artistools/nonthermal/solvespencerfanocmd.py index bd99d75f5..c113b6810 100755 --- a/artistools/nonthermal/solvespencerfanocmd.py +++ b/artistools/nonthermal/solvespencerfanocmd.py @@ -157,7 +157,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None = print("A time or timestep must be specified.") sys.exit() - modeldata, _ = at.inputmodel.get_modeldata(modelpath) + modeldata, _ = at.inputmodel.get_modeldata_pandas(modelpath) if args.velocity >= 0.0: args.modelgridindex = at.inputmodel.get_mgi_of_velocity_kms(modelpath, args.velocity) else: diff --git a/artistools/packets/packetsplots.py b/artistools/packets/packetsplots.py index 2b00c5a7f..d058dbc0b 100644 --- a/artistools/packets/packetsplots.py +++ b/artistools/packets/packetsplots.py @@ -140,7 +140,7 @@ def plot_last_emission_velocities_histogram( nrows=1, ncols=1, figsize=(5, 4), tight_layout={"pad": 1.0, "w_pad": 0.0, "h_pad": 0.5}, sharex=True ) - dfmodel, modelmeta = at.get_modeldata_polars(modelpath=modelpath, getheadersonly=True, printwarningsonly=True) + dfmodel, modelmeta = at.get_modeldata(modelpath=modelpath, getheadersonly=True, printwarningsonly=True) nprocs_read, dfpackets = at.packets.get_packets_pl( modelpath, maxpacketfiles=maxpacketfiles, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT" diff --git a/artistools/plotspherical.py b/artistools/plotspherical.py index a124edc34..fa99ae1cc 100755 --- a/artistools/plotspherical.py +++ b/artistools/plotspherical.py @@ -314,7 +314,7 @@ def main(args: argparse.Namespace | None = None, argsraw: list[str] | None = Non assert args.atomic_number is None args.atomic_number = at.get_atomic_number(args.elem) - dfmodel, modelmeta = at.get_modeldata_polars(modelpath=args.modelpath, getheadersonly=True, printwarningsonly=True) + dfmodel, modelmeta = at.get_modeldata(modelpath=args.modelpath, getheadersonly=True, printwarningsonly=True) dfestimators = at.estimators.scan_estimators(modelpath=args.modelpath) if "temperature" in args.plotvars else None nprocs_read, dfpackets = at.packets.get_packets_pl( diff --git a/artistools/spectra/plotspectra.py b/artistools/spectra/plotspectra.py index d57da7eb3..3e8d2099d 100755 --- a/artistools/spectra/plotspectra.py +++ b/artistools/spectra/plotspectra.py @@ -27,7 +27,7 @@ import artistools.spectra as atspectra from artistools.configuration import get_config -from artistools.inputmodel import get_modeldata +from artistools.inputmodel import get_modeldata_pandas from artistools.misc import AppendPath from artistools.misc import CustomArgHelpFormatter from artistools.misc import flatten_list @@ -888,10 +888,10 @@ def make_contrib_plot( import artistools.estimators as atestimators import artistools.packets as atpackets - modeldata, _ = get_modeldata(modelpath) + modeldata, _ = get_modeldata_pandas(modelpath) if args.classicartis: - modeldata, _ = get_modeldata(modelpath) + modeldata, _ = get_modeldata_pandas(modelpath) estimators = atestimators.estimators_classic.read_classic_estimators(modelpath, modeldata) allnonemptymgilist = list(modeldata.index) diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py index a390a2e54..ee1007663 100644 --- a/artistools/spectra/spectra.py +++ b/artistools/spectra/spectra.py @@ -253,9 +253,9 @@ def get_from_packets( dfpackets = dfpackets.filter(pl.col("t_arrive_d").is_between(timelowdays, timehighdays)) elif use_time == "escape": - from artistools.inputmodel import get_modeldata + from artistools.inputmodel import get_modeldata_pandas - modeldata, _ = get_modeldata(modelpath) + modeldata, _ = get_modeldata_pandas(modelpath) vmax_beta = modeldata.iloc[-1].vel_r_max_kmps * 299792.458 escapesurfacegamma = math.sqrt(1 - vmax_beta**2) diff --git a/artistools/transitions.py b/artistools/transitions.py index 2e88f237c..fe4e7e079 100644 --- a/artistools/transitions.py +++ b/artistools/transitions.py @@ -279,7 +279,7 @@ def main(args: argparse.Namespace | None = None, argsraw: Sequence[str] | None = timestep = at.get_timestep_of_timedays(modelpath, args.timedays) if args.timedays else args.timestep - modeldata, _ = at.inputmodel.get_modeldata(Path(modelpath, "model.txt")) + modeldata, _ = at.inputmodel.get_modeldata_pandas(Path(modelpath, "model.txt")) estimators_all = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex) if not estimators_all: print("no estimators") diff --git a/artistools/viewing_angles_visualization.py b/artistools/viewing_angles_visualization.py index 2f438989e..87e77e263 100755 --- a/artistools/viewing_angles_visualization.py +++ b/artistools/viewing_angles_visualization.py @@ -102,7 +102,7 @@ def viewing_angles_visualisation( sys.exit() # Load model contents - dfmodel, _modelmeta = at.get_modeldata(modelfile, derived_cols=["pos_mid"]) + dfmodel, _modelmeta = at.get_modeldata_pandas(modelfile, derived_cols=["pos_mid"]) x, y, z = (dfmodel[f"pos_{ax}_mid"].to_numpy() for ax in ("x", "y", "z")) rho = dfmodel["rho"].to_numpy() diff --git a/artistools/writecomparisondata.py b/artistools/writecomparisondata.py index 6653552fa..f594033fc 100755 --- a/artistools/writecomparisondata.py +++ b/artistools/writecomparisondata.py @@ -50,14 +50,14 @@ def write_spectra(modelpath: str | Path, selected_timesteps: Sequence[int], outf def write_ntimes_nvel(outfile: TextIOWrapper, selected_timesteps: Sequence[int], modelpath: str | Path) -> None: times = at.get_timestep_times(modelpath) - _, modelmeta = at.inputmodel.get_modeldata(modelpath, getheadersonly=True) + _, modelmeta = at.inputmodel.get_modeldata_pandas(modelpath, getheadersonly=True) outfile.write(f"#NTIMES: {len(selected_timesteps)}\n") outfile.write(f"#NVEL: {modelmeta['npts_model']}\n") outfile.write(f"#TIMES[d]: {' '.join([f'{times[ts]:.2f}' for ts in selected_timesteps])}\n") def write_single_estimator(modelpath, selected_timesteps, estimators, allnonemptymgilist, outfile, keyname) -> None: - modeldata, _modelmeta = at.inputmodel.get_modeldata(modelpath, derived_cols=["vel_r_min_kmps"]) + modeldata, _modelmeta = at.inputmodel.get_modeldata_pandas(modelpath, derived_cols=["vel_r_min_kmps"]) with Path(outfile).open("w", encoding="utf-8") as f: write_ntimes_nvel(f, selected_timesteps, modelpath) if keyname == "total_dep": @@ -91,7 +91,7 @@ def write_ionfracts( outputpath, ) -> None: times = at.get_timestep_times(modelpath) - modeldata, _modelmeta = at.inputmodel.get_modeldata(modelpath, derived_cols=["vel_r_min_kmps"]) + modeldata, _modelmeta = at.inputmodel.get_modeldata_pandas(modelpath, derived_cols=["vel_r_min_kmps"]) elementlist = at.get_composition_data(modelpath) nelements = len(elementlist) for element in range(nelements): @@ -131,7 +131,7 @@ def write_ionfracts( def write_phys(modelpath, model_id, selected_timesteps, estimators, allnonemptymgilist, outputpath) -> None: times = at.get_timestep_times(modelpath) - modeldata, modelmeta = at.inputmodel.get_modeldata(modelpath, derived_cols=["vel_r_min_kmps"]) + modeldata, modelmeta = at.inputmodel.get_modeldata_pandas(modelpath, derived_cols=["vel_r_min_kmps"]) with Path(outputpath, f"phys_{model_id}_artisnebular.txt").open("w", encoding="utf-8") as f: f.write(f"#NTIMES: {len(selected_timesteps)}\n") f.write(f"#TIMES[d]: {' '.join([f'{times[ts]:.2f}' for ts in selected_timesteps])}\n") diff --git a/pyproject.toml b/pyproject.toml index 615ee8cc1..59d37efee 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,7 @@ dependencies = [ "rich>=13.9.4", "scipy>=1.14.1", "tqdm>=4.67.1", + "typing_extensions>=4.12.2", "zstandard>=0.23.0", ]