From f32d7cb5849f5f7ad32497e683cc26f79f232a3e Mon Sep 17 00:00:00 2001 From: Luke Shingles Date: Wed, 19 Apr 2023 13:06:59 +0100 Subject: [PATCH] Get direction-binned spectra and light curves from packets files (#63) Main goal was to add --frompackets support for direction-binned spectra and light curves. It sort of became my development branch: - Sped up 3D model reading by writing .parquet binary copies of the model.txt and abundances.txt - Fixed linting errors and warnings to improve code quality - Major changes to plotartisinitialabundances command (@ccollins22 please check that this still does what you wrote this for) - Add support for zstandard and lz4 compression - fix vspecpol plotting and add a test for it --- .github/workflows/linter.yml | 7 +- .github/workflows/pytest.yml | 5 +- .gitignore | 4 +- .pre-commit-config.yaml | 29 +- artistools/__init__.py | 154 ++-- artistools/__main__.py | 21 +- artistools/atomic/__init__.py | 2 - artistools/atomic/_atomic_core.py | 18 +- artistools/codecomparison.py | 48 +- artistools/commands.py | 17 +- artistools/configuration.py | 16 +- artistools/data/__init__.py | 0 artistools/data/splitmetadata.py | 2 +- artistools/deposition.py | 2 - artistools/diskcachedecorator.py | 179 ---- artistools/estimators/__init__.py | 35 +- artistools/estimators/__main__.py | 10 +- artistools/estimators/estimators.py | 117 ++- artistools/estimators/estimators_classic.py | 97 ++- artistools/estimators/exportmassfractions.py | 9 +- .../estimators/plot3destimators_classic.py | 22 +- artistools/estimators/plotestimators.py | 81 +- artistools/gsinetwork.py | 25 +- artistools/hesma_scripts.py | 24 +- artistools/initial_composition.py | 261 +++--- artistools/inputmodel/1dslicefrom3d.py | 27 +- artistools/inputmodel/__init__.py | 34 +- artistools/inputmodel/botyanski2017.py | 7 +- artistools/inputmodel/describeinputmodel.py | 11 +- artistools/inputmodel/downscale3dgrid.py | 4 +- artistools/inputmodel/energyinputfiles.py | 36 +- artistools/inputmodel/fromcmfgen/__init__.py | 0 .../convert_to_artis_neartimezero.py | 5 +- artistools/inputmodel/fromcmfgen/rd_cmfgen.py | 16 +- artistools/inputmodel/fullymixed.py | 16 +- artistools/inputmodel/inputmodel_misc.py | 510 ++++++++---- artistools/inputmodel/lapuente.py | 9 +- artistools/inputmodel/makeartismodel.py | 4 - .../inputmodel/maketardismodelfromartis.py | 4 +- artistools/inputmodel/map_1d_to_3d_grid.py | 19 +- artistools/inputmodel/maptogrid.py | 310 ++++--- artistools/inputmodel/modelfromhydro.py | 70 +- artistools/inputmodel/plotdensity.py | 3 - artistools/inputmodel/recombinationenergy.py | 29 +- .../inputmodel/rprocess_from_trajectory.py | 114 +-- artistools/inputmodel/rprocess_solar.py | 23 +- artistools/inputmodel/scalevelocity.py | 6 +- artistools/inputmodel/shen2018.py | 6 +- .../inputmodel/slice1Dfromconein3dmodel.py | 11 +- artistools/inputmodel/test_inputmodel.py | 12 - artistools/lightcurve/__init__.py | 70 +- artistools/lightcurve/__main__.py | 10 +- artistools/lightcurve/lightcurve.py | 343 ++++---- artistools/lightcurve/plotlightcurve.py | 666 ++++++++------- artistools/lightcurve/test_lightcurve.py | 73 ++ artistools/lightcurve/viewingangleanalysis.py | 145 ++-- .../lightcurve/writebollightcurvedata.py | 17 +- artistools/linefluxes.py | 89 +- artistools/logfiles.py | 5 +- artistools/macroatom.py | 41 +- artistools/misc.py | 663 +++++++-------- artistools/nltepops/__init__.py | 20 +- artistools/nltepops/__main__.py | 11 +- artistools/nltepops/nltepops.py | 39 +- artistools/nltepops/plotnltepops.py | 96 +-- artistools/nonthermal/__init__.py | 86 +- artistools/nonthermal/__main__.py | 9 + artistools/nonthermal/_nonthermal_core.py | 127 ++- artistools/nonthermal/plotnonthermal.py | 36 +- artistools/nonthermal/solvespencerfanocmd.py | 29 +- artistools/packets/__init__.py | 551 +------------ artistools/packets/packets.py | 775 ++++++++++++++++++ artistools/packets/packetsplots.py | 15 +- artistools/plottools.py | 17 +- artistools/radfield.py | 101 ++- artistools/spectra/__init__.py | 49 +- artistools/spectra/__main__.py | 10 +- artistools/spectra/plotspectra.py | 426 +++++----- .../spectra/sampleblackbodyfrompacketTR.py | 10 +- artistools/spectra/spectra.py | 773 ++++++++--------- artistools/spectra/test_spectra.py | 56 +- artistools/spectra/test_vspectra.py | 44 + artistools/stats.py | 2 +- artistools/test_artistools.py | 54 +- artistools/transitions.py | 97 +-- artistools/viewing_angles_visualization.py | 20 +- artistools/writecomparisondata.py | 17 +- pyproject.toml | 104 ++- requirements.txt | 39 +- setup.py | 5 +- tests/data/.gitignore | 3 +- tests/data/setuptestdata.sh | 3 + 92 files changed, 4253 insertions(+), 3964 deletions(-) create mode 100644 artistools/data/__init__.py mode change 100644 => 100755 artistools/data/splitmetadata.py delete mode 100644 artistools/diskcachedecorator.py mode change 100644 => 100755 artistools/estimators/plotestimators.py mode change 100644 => 100755 artistools/initial_composition.py mode change 100644 => 100755 artistools/inputmodel/describeinputmodel.py create mode 100644 artistools/inputmodel/fromcmfgen/__init__.py mode change 100644 => 100755 artistools/inputmodel/makeartismodel.py mode change 100644 => 100755 artistools/inputmodel/maketardismodelfromartis.py mode change 100644 => 100755 artistools/inputmodel/modelfromhydro.py mode change 100644 => 100755 artistools/inputmodel/rprocess_from_trajectory.py mode change 100644 => 100755 artistools/inputmodel/rprocess_solar.py create mode 100755 artistools/lightcurve/test_lightcurve.py mode change 100644 => 100755 artistools/linefluxes.py mode change 100644 => 100755 artistools/logfiles.py mode change 100644 => 100755 artistools/macroatom.py create mode 100644 artistools/nonthermal/__main__.py mode change 100644 => 100755 artistools/nonthermal/solvespencerfanocmd.py create mode 100644 artistools/packets/packets.py mode change 100644 => 100755 artistools/radfield.py create mode 100755 artistools/spectra/test_vspectra.py mode change 100644 => 100755 artistools/transitions.py diff --git a/.github/workflows/linter.yml b/.github/workflows/linter.yml index df8f82aa9..cbb0aa84b 100644 --- a/.github/workflows/linter.yml +++ b/.github/workflows/linter.yml @@ -78,6 +78,11 @@ jobs: - name: Lint with pylint run: | pylint artistools + ruff: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: chartboost/ruff-action@v1 superlinter: name: Super linter @@ -91,7 +96,7 @@ jobs: fetch-depth: 0 - name: Lint Code Base - uses: github/super-linter/slim@v4 + uses: github/super-linter/slim@v5.0.0 env: LINTER_RULES_PATH: ./ # LOG_LEVEL: WARNING diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index c8ae6871f..8f68044f4 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -41,10 +41,9 @@ jobs: - name: Cache test data uses: actions/cache@v3 - id: cache-atomicdata with: - path: tests/data/testmodel.tar.xz - key: https://theory.gsi.de/~lshingle/artis_http_public/artistools/testmodel.tar.xz + path: tests/data/*.tar.xz + key: testdata20230417 - name: Download/extract test data working-directory: tests/data/ diff --git a/.gitignore b/.gitignore index 7748e72e2..eed5d2afe 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ __pycache__/ *.py[cod] *$py.class coverage.xml +.*_cache/ *.so @@ -61,4 +62,5 @@ plottingscripts/ _version.py -.vscode \ No newline at end of file +.vscode +.dmypy.json diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index da801b9dd..dab11db99 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,10 +7,12 @@ repos: args: [--maxkb=800] - id: check-ast - id: check-case-conflict + - id: check-builtin-literals - id: check-docstring-first - id: check-executables-have-shebangs - id: check-json - id: check-merge-conflict + - id: check-symlinks - id: check-toml - id: check-yaml - id: detect-private-key @@ -26,6 +28,11 @@ repos: rev: 0.2.2 hooks: - id: yamlfmt + - repo: https://github.com/charliermarsh/ruff-pre-commit + rev: v0.0.261 + hooks: + - id: ruff + args: [--fix, --exit-non-zero-on-fix] - repo: https://github.com/asottile/reorder_python_imports rev: v3.9.0 hooks: @@ -33,31 +40,13 @@ repos: types: [python] args: [--py39-plus, --exit-zero-even-if-changed] - repo: https://github.com/psf/black - rev: 23.1.0 + rev: 23.3.0 hooks: - id: black - args: [--quiet] - - repo: https://github.com/pycqa/flake8 - rev: 6.0.0 - hooks: - - id: flake8 - types: [python] - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.1.1 + rev: v1.2.0 hooks: - id: mypy - # language: system additional_dependencies: [numpy, types-PyYAML, types-psutil] types: [python] require_serial: true - # verbose: true - - repo: local - hooks: - - id: pylint - name: pylint - entry: pylint - language: python - # language: system - types: [python] - require_serial: true - args: [-rn, -sn, --errors-only] diff --git a/artistools/__init__.py b/artistools/__init__.py index 9936f7e99..ec9e35236 100644 --- a/artistools/__init__.py +++ b/artistools/__init__.py @@ -1,12 +1,8 @@ -#!/usr/bin/env python3 -"""Artistools. +"""artistools. A collection of plotting, analysis, and file format conversion tools for the ARTIS radiative transfer code. """ -from artistools.diskcachedecorator import diskcache # noreorder -from artistools.configuration import get_config -from artistools.configuration import set_config import artistools.atomic import artistools.codecomparison import artistools.commands @@ -18,77 +14,83 @@ import artistools.nltepops import artistools.nonthermal import artistools.packets -import artistools.plottools import artistools.radfield import artistools.spectra import artistools.transitions -from artistools.__main__ import addargs -from artistools.__main__ import main -from artistools.inputmodel import add_derived_cols_to_modeldata -from artistools.inputmodel import get_2d_modeldata -from artistools.inputmodel import get_cell_angle -from artistools.inputmodel import get_dfmodel_dimensions -from artistools.inputmodel import get_mean_cell_properties_of_angle_bin -from artistools.inputmodel import get_mgi_of_velocity_kms -from artistools.inputmodel import get_modeldata -from artistools.inputmodel import get_modeldata_tuple -from artistools.inputmodel import save_initialabundances -from artistools.inputmodel import save_modeldata -from artistools.misc import AppendPath -from artistools.misc import CustomArgHelpFormatter -from artistools.misc import decode_roman_numeral -from artistools.misc import anyexist -from artistools.misc import firstexisting -from artistools.misc import flatten_list -from artistools.misc import gather_res_data -from artistools.misc import get_artis_constants -from artistools.misc import get_atomic_number -from artistools.misc import get_bflist -from artistools.misc import get_cellsofmpirank -from artistools.misc import get_composition_data -from artistools.misc import get_composition_data_from_outputfile -from artistools.misc import get_deposition -from artistools.misc import get_elsymbol -from artistools.misc import get_elsymbolslist -from artistools.misc import get_escaped_arrivalrange -from artistools.misc import get_filterfunc -from artistools.misc import get_grid_mapping -from artistools.misc import get_inputparams -from artistools.misc import get_ionstring -from artistools.misc import get_linelist_dict -from artistools.misc import get_linelist_dataframe -from artistools.misc import read_linestatfile -from artistools.misc import get_model_name -from artistools.misc import get_mpiranklist -from artistools.misc import get_mpirankofcell -from artistools.misc import get_nprocs -from artistools.misc import get_runfolders -from artistools.misc import linetuple -from artistools.misc import get_syn_dir -from artistools.misc import get_time_range -from artistools.misc import get_timestep_of_timedays -from artistools.misc import get_timestep_time -from artistools.misc import get_timestep_times_float -from artistools.misc import get_viewinganglebin_definitions -from artistools.misc import get_vpkt_config -from artistools.misc import get_wid_init_at_tmin -from artistools.misc import get_viewingdirectionbincount -from artistools.misc import get_viewingdirection_phibincount -from artistools.misc import get_viewingdirection_costhetabincount -from artistools.misc import get_wid_init_at_tmodel -from artistools.misc import get_z_a_nucname -from artistools.misc import join_pdf_files -from artistools.misc import make_namedtuple -from artistools.misc import makelist -from artistools.misc import match_closest_time -from artistools.misc import namedtuple -from artistools.misc import parse_cdefines -from artistools.misc import parse_range -from artistools.misc import parse_range_list -from artistools.misc import readnoncommentline -from artistools.misc import roman_numerals -from artistools.misc import showtimesteptimes -from artistools.misc import stripallsuffixes -from artistools.misc import trim_or_pad -from artistools.misc import vec_len -from artistools.misc import zopen +from .__main__ import addargs +from .__main__ import main +from .configuration import get_config +from .configuration import set_config +from .inputmodel import add_derived_cols_to_modeldata +from .inputmodel import get_2d_modeldata +from .inputmodel import get_cell_angle +from .inputmodel import get_dfmodel_dimensions +from .inputmodel import get_mean_cell_properties_of_angle_bin +from .inputmodel import get_mgi_of_velocity_kms +from .inputmodel import get_modeldata +from .inputmodel import get_modeldata_tuple +from .inputmodel import save_initelemabundances +from .inputmodel import save_modeldata +from .misc import anyexist +from .misc import AppendPath +from .misc import average_direction_bins +from .misc import CustomArgHelpFormatter +from .misc import decode_roman_numeral +from .misc import firstexisting +from .misc import flatten_list +from .misc import get_atomic_number +from .misc import get_bflist +from .misc import get_cellsofmpirank +from .misc import get_composition_data +from .misc import get_composition_data_from_outputfile +from .misc import get_costhetabin_phibin_labels +from .misc import get_deposition +from .misc import get_dirbin_labels +from .misc import get_elsymbol +from .misc import get_elsymbolslist +from .misc import get_escaped_arrivalrange +from .misc import get_file_metadata +from .misc import get_filterfunc +from .misc import get_grid_mapping +from .misc import get_inputparams +from .misc import get_ionstring +from .misc import get_linelist_dataframe +from .misc import get_linelist_dict +from .misc import get_model_name +from .misc import get_mpiranklist +from .misc import get_mpirankofcell +from .misc import get_nprocs +from .misc import get_nu_grid +from .misc import get_runfolders +from .misc import get_syn_dir +from .misc import get_time_range +from .misc import get_timestep_of_timedays +from .misc import get_timestep_time +from .misc import get_timestep_times_float +from .misc import get_viewingdirection_costhetabincount +from .misc import get_viewingdirection_phibincount +from .misc import get_viewingdirectionbincount +from .misc import get_vpkt_config +from .misc import get_vspec_dir_labels +from .misc import get_wid_init_at_tmin +from .misc import get_wid_init_at_tmodel +from .misc import get_z_a_nucname +from .misc import join_pdf_files +from .misc import linetuple +from .misc import makelist +from .misc import match_closest_time +from .misc import namedtuple +from .misc import parse_range +from .misc import parse_range_list +from .misc import read_linestatfile +from .misc import readnoncommentline +from .misc import roman_numerals +from .misc import showtimesteptimes +from .misc import split_dataframe_dirbins +from .misc import stripallsuffixes +from .misc import trim_or_pad +from .misc import vec_len +from .misc import zopen +from .plottools import set_mpl_style + +set_mpl_style() diff --git a/artistools/__main__.py b/artistools/__main__.py index 40605e3f5..7ebe5d042 100644 --- a/artistools/__main__.py +++ b/artistools/__main__.py @@ -1,9 +1,7 @@ -#!/usr/bin/env python3 # PYTHON_ARGCOMPLETE_OK import argparse import importlib import multiprocessing -from typing import Union import argcomplete @@ -16,6 +14,11 @@ def addargs(parser=None) -> None: "inputmodel": { "describe": ("inputmodel.describeinputmodel", "main"), "maptogrid": ("inputmodel.maptogrid", "main"), + "makeartismodelfromparticlegridmap": ("inputmodel.modelfromhydro", "main"), + "makeartismodel": ("inputmodel.makeartismodel", "main"), + }, + "estimators": { + "plot": ("estimators.plotestimators", "main"), }, "lightcurves": { "plot": ("lightcurve.plotlightcurve", "main"), @@ -23,21 +26,21 @@ def addargs(parser=None) -> None: "spectra": { "plot": ("spectra.plotspectra", "main"), }, + "comparetogsinetwork": ("gsinetwork", "main"), } def addsubparsers(parser, parentcommand, dictcommands, depth: int = 1) -> None: - subparsers = parser.add_subparsers(dest=f"{parentcommand} command", required=True, title="test") + subparsers = parser.add_subparsers(dest=f"{parentcommand} command", required=True) for subcommand, subcommands in dictcommands.items(): + subparser = subparsers.add_parser(subcommand, help=subcommand) if isinstance(subcommands, dict): - subparser = subparsers.add_parser(subcommand) - addsubparsers(subparser, subcommand, subcommands, depth=depth + 1) else: - command = subcommand submodulename, funcname = subcommands - submodule = importlib.import_module(f"artistools.{submodulename}", package="artistools") - subparser = subparsers.add_parser(command) + submodule = importlib.import_module( + f"artistools.{submodulename.removeprefix('artistools.')}", package="artistools" + ) submodule.addargs(subparser) subparser.set_defaults(func=getattr(submodule, funcname)) @@ -45,8 +48,6 @@ def addsubparsers(parser, parentcommand, dictcommands, depth: int = 1) -> None: def main(args=None, argsraw=None, **kwargs) -> None: """Parse and run artistools commands.""" - import artistools.commands - parser = argparse.ArgumentParser() parser.set_defaults(func=None) diff --git a/artistools/atomic/__init__.py b/artistools/atomic/__init__.py index c6ad0a35d..771485979 100644 --- a/artistools/atomic/__init__.py +++ b/artistools/atomic/__init__.py @@ -1,4 +1,2 @@ -#!/usr/bin/env python3 -# import artistools as at from artistools.atomic._atomic_core import get_ionrecombratecalibration from artistools.atomic._atomic_core import get_levels diff --git a/artistools/atomic/_atomic_core.py b/artistools/atomic/_atomic_core.py index 043ce50d3..8c47456c7 100644 --- a/artistools/atomic/_atomic_core.py +++ b/artistools/atomic/_atomic_core.py @@ -130,7 +130,7 @@ def get_levels(modelpath, ionlist=None, get_transitions=False, get_photoionisati transition_filename = Path(modelpath, "transitiondata.txt") if not quiet: print(f"Reading {transition_filename.relative_to(Path(modelpath).parent)}") - with at.zopen(transition_filename, "rt") as ftransitions: + with at.zopen(transition_filename) as ftransitions: transitionsdict = { (Z, ionstage): dftransitions for Z, ionstage, dftransitions in parse_transitiondata(ftransitions, ionlist) @@ -142,11 +142,11 @@ def get_levels(modelpath, ionlist=None, get_transitions=False, get_photoionisati if not quiet: print(f"Reading {phixs_filename.relative_to(Path(modelpath).parent)}") - with at.zopen(phixs_filename, "rt") as fphixs: + with at.zopen(phixs_filename) as fphixs: for ( Z, - upperionstage, - upperionlevel, + _upperionstage, + _upperionlevel, lowerionstage, lowerionlevel, phixstargetlist, @@ -157,7 +157,7 @@ def get_levels(modelpath, ionlist=None, get_transitions=False, get_photoionisati level_lists = [] iontuple = namedtuple("ion", "Z ion_stage level_count ion_pot levels transitions") - with at.zopen(adatafilename, "rt") as fadata: + with at.zopen(adatafilename) as fadata: if not quiet: print(f"Reading {adatafilename.relative_to(Path(modelpath).parent)}") @@ -171,12 +171,12 @@ def get_levels(modelpath, ionlist=None, get_transitions=False, get_photoionisati def parse_recombratefile(frecomb): for line in frecomb: - Z, upper_ionstage, t_count = [int(x) for x in line.split()] + Z, upper_ionstage, t_count = (int(x) for x in line.split()) arr_log10t = [] arr_rrc_low_n = [] arr_rrc_total = [] for _ in range(int(t_count)): - log10t, rrc_low_n, rrc_total = [float(x) for x in frecomb.readline().split()] + log10t, rrc_low_n, rrc_total = (float(x) for x in frecomb.readline().split()) arr_log10t.append(log10t) arr_rrc_low_n.append(rrc_low_n) @@ -186,7 +186,7 @@ def parse_recombratefile(frecomb): {"log10T_e": arr_log10t, "rrc_low_n": arr_rrc_low_n, "rrc_total": arr_rrc_total} ) - recombdata_thision.eval("T_e = 10 ** log10T_e", inplace=True) + recombdata_thision = recombdata_thision.eval("T_e = 10 ** log10T_e") yield Z, upper_ionstage, recombdata_thision @@ -195,7 +195,7 @@ def parse_recombratefile(frecomb): def get_ionrecombratecalibration(modelpath): """Read recombrates file.""" recombdata = {} - with open(Path(modelpath, "recombrates.txt"), "r") as frecomb: + with Path(modelpath, "recombrates.txt").open("r") as frecomb: for Z, upper_ionstage, dfrrc in parse_recombratefile(frecomb): recombdata[(Z, upper_ionstage)] = dfrrc diff --git a/artistools/codecomparison.py b/artistools/codecomparison.py index 58dd77342..010eab60d 100644 --- a/artistools/codecomparison.py +++ b/artistools/codecomparison.py @@ -6,13 +6,13 @@ e.g., codecomparison/DDC10/artisnebular """ import math -from collections.abc import Iterable from collections.abc import Sequence from pathlib import Path from typing import Any from typing import Literal from typing import Union +import matplotlib.axes import numpy as np import pandas as pd @@ -27,8 +27,8 @@ def get_timestep_times_float( filepath = Path(at.get_config()["codecomparisondata1path"], modelname, f"phys_{modelname}_{codename}.txt") - with open(filepath, "r") as fphys: - ntimes = int(fphys.readline().replace("#NTIMES:", "")) + with open(filepath, encoding="utf-8") as fphys: + _ = int(fphys.readline().replace("#NTIMES:", "")) tmids = np.array([float(x) for x in fphys.readline().replace("#TIMES[d]:", "").split()]) tstarts = np.zeros_like(tmids) @@ -41,32 +41,35 @@ def get_timestep_times_float( if loc == "mid": return tmids - elif loc == "start": + if loc == "start": return tstarts - elif loc == "end": + if loc == "end": return tends - elif loc == "delta": + if loc == "delta": tdeltas = tends - tstarts return tdeltas - else: - raise ValueError("loc must be one of 'mid', 'start', 'end', or 'delta'") + + raise ValueError("loc must be one of 'mid', 'start', 'end', or 'delta'") -def read_reference_estimators(modelpath, modelgridindex=None, timestep=None): +def read_reference_estimators( + modelpath: Union[str, Path], + modelgridindex: Union[None, int, Sequence[int]] = None, + timestep: Union[None, int, Sequence[int]] = None, +) -> dict[tuple[int, int], Any]: """Read estimators from code comparison workshop file.""" - virtualfolder, inputmodel, codename = modelpath.parts + virtualfolder, inputmodel, codename = Path(modelpath).parts assert virtualfolder == "codecomparison" inputmodelfolder = Path(at.get_config()["codecomparisondata1path"], inputmodel) physfilepath = Path(inputmodelfolder, f"phys_{inputmodel}_{codename}.txt") - estimators = {} - cell_vel = {} + estimators: dict[tuple[int, int], Any] = {} cur_timestep = -1 cur_modelgridindex = -1 - with open(physfilepath, "r") as fphys: + with open(physfilepath) as fphys: ntimes = int(fphys.readline().replace("#NTIMES:", "")) arr_timedays = np.array([float(x) for x in fphys.readline().replace("#TIMES[d]:", "").split()]) assert len(arr_timedays) == ntimes @@ -81,12 +84,13 @@ def read_reference_estimators(modelpath, modelgridindex=None, timestep=None): assert np.isclose(timedays, arr_timedays[cur_timestep], rtol=0.01) elif row[0] == "#NVEL:": - nvel = int(row[1]) + _ = int(row[1]) elif not line.lstrip().startswith("#"): cur_modelgridindex += 1 key = (cur_timestep, cur_modelgridindex) + if key not in estimators: estimators[key] = {"emptycell": False} @@ -102,7 +106,7 @@ def read_reference_estimators(modelpath, modelgridindex=None, timestep=None): for ionfracfilepath in ionfracfilepaths: _, element, _, _ = ionfracfilepath.stem.split("_") - with open(ionfracfilepath, "r") as fions: + with open(ionfracfilepath) as fions: print(ionfracfilepath) ntimes_2 = int(fions.readline().replace("#NTIMES:", "")) assert ntimes_2 == ntimes @@ -172,7 +176,7 @@ def read_reference_estimators(modelpath, modelgridindex=None, timestep=None): return estimators -def get_spectra(modelpath): +def get_spectra(modelpath: Union[str, Path]) -> tuple[pd.DataFrame, np.ndarray]: modelpath = Path(modelpath) virtualfolder, inputmodel, codename = modelpath.parts assert virtualfolder == "codecomparison" @@ -181,20 +185,22 @@ def get_spectra(modelpath): specfilepath = Path(inputmodelfolder, f"spectra_{inputmodel}_{codename}.txt") - with open(specfilepath, "r") as fspec: + with open(specfilepath) as fspec: ntimes = int(fspec.readline().replace("#NTIMES:", "")) - nwave = int(fspec.readline().replace("#NWAVE:", "")) + _ = int(fspec.readline().replace("#NWAVE:", "")) arr_timedays = np.array([float(x) for x in fspec.readline().split()[1:]]) assert len(arr_timedays) == ntimes dfspectra = pd.read_csv( - fspec, delim_whitespace=True, header=None, names=["lambda"] + list(arr_timedays), comment="#" + fspec, delim_whitespace=True, header=None, names=["lambda", *list(arr_timedays)], comment="#" ) return dfspectra, arr_timedays -def plot_spectrum(modelpath, timedays, ax, **plotkwargs): +def plot_spectrum( + modelpath: Union[str, Path], timedays: Union[str, float], axis: matplotlib.axes.Axes, **plotkwargs +) -> None: dfspectra, arr_timedays = get_spectra(modelpath) # print(dfspectra) timeindex = (np.abs(arr_timedays - float(timedays))).argmin() @@ -209,4 +215,4 @@ def plot_spectrum(modelpath, timedays, ax, **plotkwargs): megaparsec_to_cm = 3.085677581491367e24 arr_flux = dfspectra[dfspectra.columns[timeindex + 1]] / 4 / math.pi / (megaparsec_to_cm**2) - ax.plot(dfspectra["lambda"], arr_flux, label=label, **plotkwargs) + axis.plot(dfspectra["lambda"], arr_flux, label=label, **plotkwargs) diff --git a/artistools/commands.py b/artistools/commands.py index 165bcf9f3..3facf426b 100644 --- a/artistools/commands.py +++ b/artistools/commands.py @@ -2,7 +2,7 @@ from pathlib import Path -def get_commandlist(): +def get_commandlist() -> dict[str, tuple[str, str]]: commandlist = { "at": ("artistools", "main"), "artistools": ("artistools", "main"), @@ -60,37 +60,36 @@ def get_commandlist(): return commandlist -def get_console_scripts(): +def get_console_scripts() -> list[str]: console_scripts = [ f"{command} = {submodulename}:{funcname}" for command, (submodulename, funcname) in get_commandlist().items() ] return console_scripts -def setup_completions(): +def setup_completions() -> None: # Add the following lines to your .zshrc file to get command completion: # autoload -U bashcompinit # bashcompinit # source artistoolscompletions.sh path_repo = Path(__file__).absolute().parent.parent - completioncommands = [] with open(path_repo / "artistoolscompletions.sh", "w", encoding="utf-8") as f: f.write("#!/usr/bin/env zsh\n") - proc = subprocess.run(["register-python-argcomplete", "__MY_COMMAND__"], capture_output=True, text=True) + proc = subprocess.run( + ["register-python-argcomplete", "__MY_COMMAND__"], capture_output=True, text=True, check=True + ) if proc.stderr: print(proc.stderr) - scriptlines = proc.stdout - strfunctiondefs, strsplit, strcommandregister = proc.stdout.rpartition("}\n") f.write(strfunctiondefs) f.write(strsplit) f.write("\n\n") - for command in get_commandlist().keys(): + for command in get_commandlist(): completecommand = strcommandregister.replace("__MY_COMMAND__", command) f.write(completecommand + "\n") @@ -98,5 +97,5 @@ def setup_completions(): print("source artistoolscompletions.sh") -def addargs(parser=None): +def addargs(parser=None) -> None: pass diff --git a/artistools/configuration.py b/artistools/configuration.py index 69276afae..6d9df2c83 100644 --- a/artistools/configuration.py +++ b/artistools/configuration.py @@ -10,8 +10,6 @@ def setup_config(): - global config - mp.set_start_method("fork") # count the cores (excluding the efficiency cores on ARM) try: @@ -32,17 +30,8 @@ def setup_config(): # print(f"Using {num_processes} processes") - config["enable_diskcache"] = False config["num_processes"] = num_processes - # pyarrow is faster, if it is installed - try: - import pyarrow - - config["pandas_engine"] = "pyarrow" - except ImportError: - config["pandas_engine"] = "c" - config["figwidth"] = 5 config["codecomparisondata1path"] = Path( "/Users/luke/Library/Mobile Documents/com~apple~CloudDocs/GitHub/sn-rad-trans/data1" @@ -58,13 +47,12 @@ def setup_config(): def get_config(key: Optional[str] = None): - global config if not config: setup_config() if key is None: return config - else: - return config[key] + + return config[key] def set_config(key: str, value: Any) -> None: diff --git a/artistools/data/__init__.py b/artistools/data/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/artistools/data/splitmetadata.py b/artistools/data/splitmetadata.py old mode 100644 new mode 100755 index 5f240ec59..cb260e245 --- a/artistools/data/splitmetadata.py +++ b/artistools/data/splitmetadata.py @@ -8,7 +8,7 @@ def main(): with Path("metadata.yml").open("r") as yamlfile: metadata = yaml.load(yamlfile, Loader=yaml.FullLoader) - for obsfile in metadata.keys(): + for obsfile in metadata: metafilepath = Path(obsfile).with_suffix(Path(obsfile).suffix + ".meta.yml") with open(metafilepath, "w") as metafile: yaml.dump(metadata[obsfile], metafile) diff --git a/artistools/deposition.py b/artistools/deposition.py index 0c8cd4e59..8017b8b68 100755 --- a/artistools/deposition.py +++ b/artistools/deposition.py @@ -7,8 +7,6 @@ from astropy import units as u import artistools as at -import artistools.inputmodel -import artistools.nltepops def forward_doubledecay( diff --git a/artistools/diskcachedecorator.py b/artistools/diskcachedecorator.py deleted file mode 100644 index 15279dea3..000000000 --- a/artistools/diskcachedecorator.py +++ /dev/null @@ -1,179 +0,0 @@ -import lzma -import os.path -import time -from functools import wraps -from pathlib import Path - -import artistools.configuration -import artistools.misc - - -def diskcache( - ignoreargs=[], ignorekwargs=[], saveonly=False, quiet=False, savezipped=False, funcdepends=None, funcversion=None -): - import pickle - import hashlib - - def printopt(*args, **kwargs): - if not quiet: - print(*args, **kwargs) - - @wraps(diskcache) - def diskcacheinner(func): - @wraps(func) - def wrapper(*args, **kwargs): - # save cached files in the folder of the first file/folder specified in the arguments - modelpath = None - for arg in [*args, *kwargs.values()]: - if modelpath is None: - try: - if os.path.isfile(arg): - modelpath = Path(arg).parent - - except TypeError: - pass - - for arg in [*args, *kwargs.values()]: - if modelpath is None: - try: - if os.path.isdir(arg): - modelpath = arg - except TypeError: - pass - - if modelpath is None: - modelpath = Path() # use current folder - - cachefolder = Path(modelpath, "__artistoolscache__.nosync") - - if cachefolder.is_dir(): - try: - import xattr - - xattr.setxattr(cachefolder, "com.dropbox.ignored", b"1") - except OSError: - pass - except ModuleNotFoundError: - pass - - namearghash = hashlib.sha1() - namearghash.update(func.__module__.encode("utf-8")) - namearghash.update(func.__qualname__.encode("utf-8")) - - namearghash.update( - str(tuple(arg for argindex, arg in enumerate(args) if argindex not in ignoreargs)).encode("utf-8") - ) - - namearghash.update(str({k: v for k, v in kwargs.items() if k not in ignorekwargs}).encode("utf-8")) - - namearghash_strhex = namearghash.hexdigest() - - # make sure modifications to any file paths in the arguments will trigger an update - argfilesmodifiedhash = hashlib.sha1() - for arg in args: - try: - if os.path.isfile(arg): - argfilesmodifiedhash.update(str(os.path.getmtime(arg)).encode("utf-8")) - except TypeError: - pass - argfilesmodifiedhash_strhex = "_filesmodifiedhash_" + argfilesmodifiedhash.hexdigest() - - filename_nogz = Path(cachefolder, f"cached-{func.__module__}.{func.__qualname__}-{namearghash_strhex}.tmp") - filename_xz = filename_nogz.with_suffix(".tmp.xz") - filename_gz = filename_nogz.with_suffix(".tmp.gz") - - execfunc = True - saveresult = False - functime = -1 - - if (filename_nogz.exists() or filename_xz.exists() or filename_gz.exists()) and not saveonly: - # found a candidate file, so load it - filename = ( - filename_nogz if filename_nogz.exists() else filename_gz if filename_gz.exists() else filename_xz - ) - - filesize = Path(filename).stat().st_size / 1024 / 1024 - - try: - printopt(f"diskcache: Loading '{filename}' ({filesize:.1f} MiB)...") - - with artistools.misc.zopen(filename, "rb") as f: - result, version_filein = pickle.load(f) - - if version_filein == str_funcversion + argfilesmodifiedhash_strhex: - execfunc = False - elif (not funcversion) and (not version_filein.startswith("funcversion_")): - execfunc = False - # elif version_filein == sourcehash_strhex: - # execfunc = False - else: - printopt(f"diskcache: Overwriting '{filename}' (function version mismatch or file modified)") - - except Exception as ex: - # ex = sys.exc_info()[0] - printopt(f"diskcache: Overwriting '{filename}' (Error: {ex})") - pass - - if execfunc: - timestart = time.perf_counter() - result = func(*args, **kwargs) - functime = time.perf_counter() - timestart - - if functime > 1: - # slow functions are worth saving to disk - saveresult = True - else: - # check if we need to replace the gzipped or non-gzipped file with the correct one - # if we so, need to save the new file even though functime is unknown since we read - # from disk version instead of executing the function - if savezipped and filename_nogz.exists(): - saveresult = True - elif not savezipped and filename_xz.exists(): - saveresult = True - - if saveresult: - # if the cache folder doesn't exist, create it - if not cachefolder.is_dir(): - cachefolder.mkdir(parents=True, exist_ok=True) - try: - import xattr - - xattr.setxattr(cachefolder, "com.dropbox.ignored", b"1") - except OSError: - pass - except ModuleNotFoundError: - pass - - if filename_nogz.exists(): - filename_nogz.unlink() - if filename_gz.exists(): - filename_gz.unlink() - if filename_xz.exists(): - filename_xz.unlink() - - fopen, filename = (lzma.open, filename_xz) if savezipped else (open, filename_nogz) - with fopen(filename, "wb") as f: - pickle.dump( - (result, str_funcversion + argfilesmodifiedhash_strhex), f, protocol=pickle.HIGHEST_PROTOCOL - ) - - filesize = Path(filename).stat().st_size / 1024 / 1024 - printopt(f"diskcache: Saved '{filename}' ({filesize:.1f} MiB, functime {functime:.1f}s)") - - return result - - # sourcehash = hashlib.sha1() - # sourcehash.update(inspect.getsource(func).encode('utf-8')) - # if funcdepends: - # try: - # for f in funcdepends: - # sourcehash.update(inspect.getsource(f).encode('utf-8')) - # except TypeError: - # sourcehash.update(inspect.getsource(funcdepends).encode('utf-8')) - # - # sourcehash_strhex = sourcehash.hexdigest() - str_funcversion = f"funcversion_{funcversion}" if funcversion else "funcversion_none" - - return wrapper if artistools.configuration.get_config()["enable_diskcache"] else func - - return diskcacheinner diff --git a/artistools/estimators/__init__.py b/artistools/estimators/__init__.py index 69f7b9a98..642237e7f 100644 --- a/artistools/estimators/__init__.py +++ b/artistools/estimators/__init__.py @@ -1,18 +1,17 @@ -#!/usr/bin/env python3 -"""Artistools - spectra related functions.""" -from artistools.estimators.estimators import apply_filters -from artistools.estimators.estimators import get_averaged_estimators -from artistools.estimators.estimators import get_averageexcitation -from artistools.estimators.estimators import get_averageionisation -from artistools.estimators.estimators import get_dictlabelreplacements -from artistools.estimators.estimators import get_ionrecombrates_fromfile -from artistools.estimators.estimators import get_partiallycompletetimesteps -from artistools.estimators.estimators import get_units_string -from artistools.estimators.estimators import get_variablelongunits -from artistools.estimators.estimators import get_variableunits -from artistools.estimators.estimators import parse_estimfile -from artistools.estimators.estimators import read_estimators -from artistools.estimators.estimators import read_estimators_from_file -from artistools.estimators.plotestimators import addargs -from artistools.estimators.plotestimators import main -from artistools.estimators.plotestimators import main as plot +"""Artistools - estimators related functions.""" +from .__main__ import main +from .estimators import apply_filters +from .estimators import get_averaged_estimators +from .estimators import get_averageexcitation +from .estimators import get_averageionisation +from .estimators import get_dictlabelreplacements +from .estimators import get_ionrecombrates_fromfile +from .estimators import get_partiallycompletetimesteps +from .estimators import get_units_string +from .estimators import get_variablelongunits +from .estimators import get_variableunits +from .estimators import parse_estimfile +from .estimators import read_estimators +from .estimators import read_estimators_from_file +from .plotestimators import addargs +from .plotestimators import main as plot diff --git a/artistools/estimators/__main__.py b/artistools/estimators/__main__.py index b6565bcf7..9ab67be8e 100644 --- a/artistools/estimators/__main__.py +++ b/artistools/estimators/__main__.py @@ -1,8 +1,12 @@ import multiprocessing -import artistools as at -import artistools.estimators.plotestimators +from .plotestimators import main as plot + + +def main() -> None: + plot() + if __name__ == "__main__": multiprocessing.freeze_support() - at.estimators.plotestimators.main() + main() diff --git a/artistools/estimators/estimators.py b/artistools/estimators/estimators.py index d30da045e..c0d3d2829 100755 --- a/artistools/estimators/estimators.py +++ b/artistools/estimators/estimators.py @@ -3,7 +3,6 @@ Examples are temperatures, populations, and heating/cooling rates. """ -# import math import argparse import math import multiprocessing @@ -26,8 +25,6 @@ import artistools as at import artistools.nltepops -# from itertools import chain - def get_variableunits(key: Optional[str] = None) -> Union[str, dict[str, str]]: variableunits = { @@ -85,7 +82,7 @@ def get_ionrecombrates_fromfile(filename: Union[Path, str]) -> pd.DataFrame: print(f"Reading {filename}") header_row = [] - with open(filename, "r") as filein: + with open(filename) as filein: while True: line = filein.readline() if line.strip().startswith("TOTAL RECOMBINATION RATE"): @@ -127,12 +124,15 @@ def get_units_string(variable: str) -> str: def parse_estimfile( - estfilepath: Path, modelpath: Path, get_ion_values: bool = True, get_heatingcooling: bool = True -) -> Iterator[tuple[int, int, dict]]: + estfilepath: Path, + modelpath: Path, + get_ion_values: bool = True, + get_heatingcooling: bool = True, +) -> Iterator[tuple[int, int, dict]]: # pylint: disable=unused-argument """Generate timestep, modelgridindex, dict from estimator file.""" # itstep = at.get_inputparams(modelpath)['itstep'] - with at.zopen(estfilepath, "rt") as estimfile: + with at.zopen(estfilepath) as estimfile: timestep: int = -1 modelgridindex: int = -1 estimblock: dict[Any, Any] = {} @@ -235,7 +235,6 @@ def parse_estimfile( yield timestep, modelgridindex, estimblock -# @at.diskcache(ignorekwargs=['printfilename'], quiet=False, funcdepends=parse_estimfile, savezipped=True) def read_estimators_from_file( folderpath: Union[Path, str], modelpath: Path, @@ -247,15 +246,12 @@ def read_estimators_from_file( ) -> dict[tuple[int, int], Any]: estimators_thisfile = {} estimfilename = f"estimators_{mpirank:04d}.out" - estfilepath = Path(folderpath, estimfilename) - if not estfilepath.is_file(): - estfilepath = Path(folderpath, estimfilename + ".gz") - if not estfilepath.is_file(): - estfilepath = Path(folderpath, estimfilename + ".xz") - if not estfilepath.is_file(): - # not worth printing and error, because ranks with no cells to update do not produce an estimator file - # print(f'Warning: Could not find {estfilepath.relative_to(modelpath.parent)}') - return {} + try: + estfilepath = at.firstexisting(estimfilename, folder=folderpath, tryzipped=True) + except FileNotFoundError: + # not worth printing an error, because ranks with no cells to update do not produce an estimator file + # print(f'Warning: Could not find {estfilepath.relative_to(modelpath.parent)}') + return {} if printfilename: filesize = Path(estfilepath).stat().st_size / 1024 / 1024 @@ -274,11 +270,10 @@ def read_estimators_from_file( @lru_cache(maxsize=16) -# @at.diskcache(savezipped=True, funcdepends=[read_estimators_from_file, parse_estimfile]) def read_estimators( modelpath: Union[Path, str], - modelgridindex: Optional[Union[int, Sequence[int]]] = None, - timestep: Optional[Union[int, Sequence[int]]] = None, + modelgridindex: Union[None, int, Sequence[int]] = None, + timestep: Union[None, int, Sequence[int]] = None, get_ion_values: bool = True, get_heatingcooling: bool = True, ) -> dict[tuple[int, int], dict]: @@ -316,7 +311,7 @@ def read_estimators( modeldata, _ = at.inputmodel.get_modeldata(modelpath, getheadersonly=True) if "velocity_outer" in modeldata.columns: modeldata, _ = at.inputmodel.get_modeldata(modelpath) - arr_velocity_outer = tuple(list([float(v) for v in modeldata["velocity_outer"].values])) + arr_velocity_outer = tuple([float(v) for v in modeldata["velocity_outer"].to_numpy()]) else: arr_velocity_outer = None @@ -348,7 +343,7 @@ def read_estimators( arr_rankestimators = [processfile(rank) for rank in mpiranklist] for mpirank, estimators_thisfile in zip(mpiranklist, arr_rankestimators): - dupekeys = list(sorted([k for k in estimators_thisfile if k in estimators])) + dupekeys = sorted([k for k in estimators_thisfile if k in estimators]) for k in dupekeys: # dropping the lowest timestep is normal for restarts. Only warn about other cases if k[0] != dupekeys[0][0]: @@ -382,28 +377,28 @@ def get_averaged_estimators( # if single timestep, no averaging needed if isinstance(timesteps, int): - return reduce(lambda d, k: d[k], [(timesteps, modelgridindex)] + keys, estimators) + return reduce(lambda d, k: d[k], [(timesteps, modelgridindex), *keys], estimators) - firsttimestepvalue = reduce(lambda d, k: d[k], [(timesteps[0], modelgridindex)] + keys, estimators) + firsttimestepvalue = reduce(lambda d, k: d[k], [(timesteps[0], modelgridindex), *keys], estimators) if isinstance(firsttimestepvalue, dict): dictout = { - k: get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, keys + [k]) - for k in firsttimestepvalue.keys() + k: get_averaged_estimators(modelpath, estimators, timesteps, modelgridindex, [*keys, k]) + for k in firsttimestepvalue } return dictout - else: - tdeltas = at.get_timestep_times_float(modelpath, loc="delta") - valuesum = 0 - tdeltasum = 0 - for timestep, tdelta in zip(timesteps, tdeltas): - for mgi in range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1): - try: - valuesum += reduce(lambda d, k: d[k], [(timestep, mgi)] + keys, estimators) * tdelta - tdeltasum += tdelta - except KeyError: - pass - return valuesum / tdeltasum + + tdeltas = at.get_timestep_times_float(modelpath, loc="delta") + valuesum = 0 + tdeltasum = 0 + for timestep, tdelta in zip(timesteps, tdeltas): + for mgi in range(modelgridindex - avgadjcells, modelgridindex + avgadjcells + 1): + try: + valuesum += reduce(lambda d, k: d[k], [(timestep, mgi), *keys], estimators) * tdelta + tdeltasum += tdelta + except KeyError: + pass + return valuesum / tdeltasum # except KeyError: # if (timestep, modelgridindex) in estimators: @@ -418,7 +413,7 @@ def get_averageionisation(populations: dict[Any, float], atomic_number: int) -> free_electron_weighted_pop_sum = 0.0 found = False popsum = 0.0 - for key in populations.keys(): + for key in populations: if isinstance(key, tuple) and key[0] == atomic_number: found = True ion_stage = key[1] @@ -442,32 +437,32 @@ def get_averageexcitation( ionpopsum = 0 if dfnltepops.empty: return float("NaN") - else: - dfnltepops_ion = dfnltepops.query( - "modelgridindex==@modelgridindex and timestep==@timestep and Z==@atomic_number & ion_stage==@ion_stage" - ) - k_b = 8.617333262145179e-05 # eV / K + dfnltepops_ion = dfnltepops.query( + "modelgridindex==@modelgridindex and timestep==@timestep and Z==@atomic_number & ion_stage==@ion_stage" + ) - ionpopsum = dfnltepops_ion.n_NLTE.sum() - energypopsum = ( - dfnltepops_ion[dfnltepops_ion.level >= 0].eval("@ionlevels.iloc[level].energy_ev.values * n_NLTE").sum() - ) + k_b = 8.617333262145179e-05 # eV / K - try: - superlevelrow = dfnltepops_ion[dfnltepops_ion.level < 0].iloc[0] - levelnumber_sl = dfnltepops_ion.level.max() + 1 + ionpopsum = dfnltepops_ion.n_NLTE.sum() + energypopsum = ( + dfnltepops_ion[dfnltepops_ion.level >= 0].eval("@ionlevels.iloc[level].energy_ev.values * n_NLTE").sum() + ) - energy_boltzfac_sum = ( - ionlevels.iloc[levelnumber_sl:].eval("energy_ev * g * exp(- energy_ev / @k_b / @T_exc)").sum() - ) + try: + superlevelrow = dfnltepops_ion[dfnltepops_ion.level < 0].iloc[0] + levelnumber_sl = dfnltepops_ion.level.max() + 1 + + energy_boltzfac_sum = ( + ionlevels.iloc[levelnumber_sl:].eval("energy_ev * g * exp(- energy_ev / @k_b / @T_exc)").sum() + ) - boltzfac_sum = ionlevels.iloc[levelnumber_sl:].eval("g * exp(- energy_ev / @k_b / @T_exc)").sum() - # adjust to the actual superlevel population from ARTIS - energypopsum += energy_boltzfac_sum * superlevelrow.n_NLTE / boltzfac_sum - except IndexError: - # no superlevel - pass + boltzfac_sum = ionlevels.iloc[levelnumber_sl:].eval("g * exp(- energy_ev / @k_b / @T_exc)").sum() + # adjust to the actual superlevel population from ARTIS + energypopsum += energy_boltzfac_sum * superlevelrow.n_NLTE / boltzfac_sum + except IndexError: + # no superlevel + pass return energypopsum / ionpopsum @@ -479,7 +474,7 @@ def get_partiallycompletetimesteps(estimators: dict[Any, Any]) -> list[int]: """ timestepcells: dict[int, list[int]] = {} all_mgis = set() - for nts, mgi in estimators.keys(): + for nts, mgi in estimators: if nts not in timestepcells: timestepcells[nts] = [] timestepcells[nts].append(mgi) diff --git a/artistools/estimators/estimators_classic.py b/artistools/estimators/estimators_classic.py index 554ffb70d..9b8fb54d2 100644 --- a/artistools/estimators/estimators_classic.py +++ b/artistools/estimators/estimators_classic.py @@ -10,7 +10,7 @@ def get_atomic_composition(modelpath): """Read ion list from output file""" atomic_composition = {} - output = open(modelpath / "output_0-0.txt", "r").read().splitlines() + output = open(modelpath / "output_0-0.txt").read().splitlines() ioncount = 0 for row in output: if row.split()[0] == "[input.c]": @@ -56,13 +56,13 @@ def get_estimator_files(modelpath): def get_first_ts_in_run_directory(modelpath): - folderlist_all = tuple(sorted([child for child in Path(modelpath).iterdir() if child.is_dir()]) + [Path(modelpath)]) + folderlist_all = (*sorted([child for child in Path(modelpath).iterdir() if child.is_dir()]), Path(modelpath)) first_timesteps_in_dir = {} for folder in folderlist_all: if os.path.isfile(folder / "output_0-0.txt"): - with open(folder / "output_0-0.txt", "r") as output_0: + with open(folder / "output_0-0.txt") as output_0: timesteps_in_dir = [ line.strip("...\n").split(" ")[-1] for line in output_0 @@ -96,57 +96,56 @@ def read_classic_estimators(modelpath, modeldata, readonly_mgi=False, readonly_t timestep = first_timesteps_in_dir[str(estfile).split("/")[0]] # get the starting timestep for the estfile # timestep = first_timesteps_in_dir[str(estfile[:-20])] # timestep = 0 # if the first timestep in the file is 0 then this is fine - with opener(estfile, "rt") as estfile: + with opener(estfile) as estfile: modelgridindex = -1 for line in estfile: row = line.split() - if int(row[0]) < int(modelgridindex): - timestep += 1 - elif int(row[0]) == int(modelgridindex): + if int(row[0]) <= int(modelgridindex): timestep += 1 modelgridindex = int(row[0]) - if readonly_mgi is False or modelgridindex in readonly_mgi: - if readonly_timestep is False or timestep in readonly_timestep: - estimators[(timestep, modelgridindex)] = {} - - if ndimensions == 1: - estimators[(timestep, modelgridindex)]["velocity_outer"] = modeldata["velocity_outer"][ - modelgridindex - ] - - estimators[(timestep, modelgridindex)]["TR"] = float(row[1]) - estimators[(timestep, modelgridindex)]["Te"] = float(row[2]) - estimators[(timestep, modelgridindex)]["W"] = float(row[3]) - estimators[(timestep, modelgridindex)]["TJ"] = float(row[4]) - - parse_ion_row_classic(row, estimators[(timestep, modelgridindex)], atomic_composition) - - # heatingrates[tid].ff, heatingrates[tid].bf, heatingrates[tid].collisional, heatingrates[tid].gamma, - # coolingrates[tid].ff, coolingrates[tid].fb, coolingrates[tid].collisional, coolingrates[tid].adiabatic) - - estimators[(timestep, modelgridindex)]["heating_ff"] = float(row[-9]) - estimators[(timestep, modelgridindex)]["heating_bf"] = float(row[-8]) - estimators[(timestep, modelgridindex)]["heating_coll"] = float(row[-7]) - estimators[(timestep, modelgridindex)]["heating_dep"] = float(row[-6]) - - estimators[(timestep, modelgridindex)]["cooling_ff"] = float(row[-5]) - estimators[(timestep, modelgridindex)]["cooling_fb"] = float(row[-4]) - estimators[(timestep, modelgridindex)]["cooling_coll"] = float(row[-3]) - estimators[(timestep, modelgridindex)]["cooling_adiabatic"] = float(row[-2]) - - # estimators[(timestep, modelgridindex)]['cooling_coll - heating_coll'] = \ - # estimators[(timestep, modelgridindex)]['cooling_coll'] - estimators[(timestep, modelgridindex)]['heating_coll'] - # - # estimators[(timestep, modelgridindex)]['cooling_fb - heating_bf'] = \ - # estimators[(timestep, modelgridindex)]['cooling_fb'] - estimators[(timestep, modelgridindex)]['heating_bf'] - # - # estimators[(timestep, modelgridindex)]['cooling_ff - heating_ff'] = \ - # estimators[(timestep, modelgridindex)]['cooling_ff'] - estimators[(timestep, modelgridindex)]['heating_ff'] - # - # estimators[(timestep, modelgridindex)]['cooling_adiabatic - heating_dep'] = \ - # estimators[(timestep, modelgridindex)]['cooling_adiabatic'] - estimators[(timestep, modelgridindex)]['heating_dep'] - - estimators[(timestep, modelgridindex)]["energy_deposition"] = float(row[-1]) + if (readonly_mgi is False or modelgridindex in readonly_mgi) and ( + readonly_timestep is False or timestep in readonly_timestep + ): + estimators[(timestep, modelgridindex)] = {} + + if ndimensions == 1: + estimators[(timestep, modelgridindex)]["velocity_outer"] = modeldata["velocity_outer"][ + modelgridindex + ] + + estimators[(timestep, modelgridindex)]["TR"] = float(row[1]) + estimators[(timestep, modelgridindex)]["Te"] = float(row[2]) + estimators[(timestep, modelgridindex)]["W"] = float(row[3]) + estimators[(timestep, modelgridindex)]["TJ"] = float(row[4]) + + parse_ion_row_classic(row, estimators[(timestep, modelgridindex)], atomic_composition) + + # heatingrates[tid].ff, heatingrates[tid].bf, heatingrates[tid].collisional, heatingrates[tid].gamma, + # coolingrates[tid].ff, coolingrates[tid].fb, coolingrates[tid].collisional, coolingrates[tid].adiabatic) + + estimators[(timestep, modelgridindex)]["heating_ff"] = float(row[-9]) + estimators[(timestep, modelgridindex)]["heating_bf"] = float(row[-8]) + estimators[(timestep, modelgridindex)]["heating_coll"] = float(row[-7]) + estimators[(timestep, modelgridindex)]["heating_dep"] = float(row[-6]) + + estimators[(timestep, modelgridindex)]["cooling_ff"] = float(row[-5]) + estimators[(timestep, modelgridindex)]["cooling_fb"] = float(row[-4]) + estimators[(timestep, modelgridindex)]["cooling_coll"] = float(row[-3]) + estimators[(timestep, modelgridindex)]["cooling_adiabatic"] = float(row[-2]) + + # estimators[(timestep, modelgridindex)]['cooling_coll - heating_coll'] = \ + # estimators[(timestep, modelgridindex)]['cooling_coll'] - estimators[(timestep, modelgridindex)]['heating_coll'] + # + # estimators[(timestep, modelgridindex)]['cooling_fb - heating_bf'] = \ + # estimators[(timestep, modelgridindex)]['cooling_fb'] - estimators[(timestep, modelgridindex)]['heating_bf'] + # + # estimators[(timestep, modelgridindex)]['cooling_ff - heating_ff'] = \ + # estimators[(timestep, modelgridindex)]['cooling_ff'] - estimators[(timestep, modelgridindex)]['heating_ff'] + # + # estimators[(timestep, modelgridindex)]['cooling_adiabatic - heating_dep'] = \ + # estimators[(timestep, modelgridindex)]['cooling_adiabatic'] - estimators[(timestep, modelgridindex)]['heating_dep'] + + estimators[(timestep, modelgridindex)]["energy_deposition"] = float(row[-1]) return estimators diff --git a/artistools/estimators/exportmassfractions.py b/artistools/estimators/exportmassfractions.py index 2bcd756c1..720c4038a 100755 --- a/artistools/estimators/exportmassfractions.py +++ b/artistools/estimators/exportmassfractions.py @@ -5,7 +5,6 @@ import numpy as np import artistools as at -import artistools.estimators def addargs(parser: argparse.ArgumentParser) -> None: @@ -26,7 +25,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: timestep = 14 elmass = {el.Z: el.mass for _, el in at.get_composition_data(modelpath).iterrows()} outfilename = args.outputpath - with open(outfilename, "wt") as fout: + with open(outfilename, "w") as fout: modelgridindexlist = range(10) estimators = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindexlist) for modelgridindex in modelgridindexlist: @@ -35,7 +34,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: numberdens = {} totaldens = 0.0 # number density times atomic mass summed over all elements - for key in popdict.keys(): + for key in popdict: try: atomic_number = int(key) numberdens[atomic_number] = popdict[atomic_number] @@ -47,12 +46,12 @@ def main(args=None, argsraw=None, **kwargs) -> None: massfracs = { atomic_number: numberdens[atomic_number] * elmass[atomic_number] / totaldens - for atomic_number in numberdens.keys() + for atomic_number in numberdens } fout.write(f"{tdays}d shell {modelgridindex}\n") massfracsum = 0.0 - for atomic_number in massfracs.keys(): + for atomic_number in massfracs: massfracsum += massfracs[atomic_number] fout.write(f"{atomic_number} {at.get_elsymbol(atomic_number)} {massfracs[atomic_number]}\n") diff --git a/artistools/estimators/plot3destimators_classic.py b/artistools/estimators/plot3destimators_classic.py index a510a11e6..4cdb6c3b2 100644 --- a/artistools/estimators/plot3destimators_classic.py +++ b/artistools/estimators/plot3destimators_classic.py @@ -6,9 +6,6 @@ import pyvista as pv import artistools as at -import artistools.estimators.estimators_classic -import artistools.initial_composition -import artistools.inputmodel.slice1Dfromconein3dmodel CLIGHT = 2.99792458e10 @@ -34,8 +31,8 @@ def get_modelgridcells_along_axis(modelpath): def get_modelgridcells_2D_slice(modeldata, modelpath): sliceaxis = "z" - slice = at.initial_composition.get_2D_slice_through_3d_model(modeldata, sliceaxis) - readonly_mgi = get_mgi_of_modeldata(slice, modelpath) + slicedata = at.initial_composition.get_2D_slice_through_3d_model(modeldata, sliceaxis) + readonly_mgi = get_mgi_of_modeldata(slicedata, modelpath) return readonly_mgi @@ -43,7 +40,7 @@ def get_modelgridcells_2D_slice(modeldata, modelpath): def get_mgi_of_modeldata(modeldata, modelpath): assoc_cells, mgi_of_propcells = at.get_grid_mapping(modelpath=modelpath) readonly_mgi = [] - for index, row in modeldata.iterrows(): + for _index, row in modeldata.iterrows(): if row["rho"] > 0: mgi = mgi_of_propcells[int(row["inputcellid"]) - 1] readonly_mgi.append(mgi) @@ -58,7 +55,7 @@ def plot_Te_vs_time_lineofsight_3d_model(modelpath, modeldata, estimators, reado associated_modeldata_row_for_mgi = modeldata.loc[modeldata["inputcellid"] == assoc_cells[mgi][0]] Te = [estimators[(timestep, mgi)]["Te"] for timestep, _ in enumerate(times)] - plt.scatter(times, Te, label=f'vel={associated_modeldata_row_for_mgi["vel_y_mid"].values[0] / CLIGHT}') + plt.scatter(times, Te, label=f'vel={associated_modeldata_row_for_mgi["vel_y_mid"].to_numpy()[0] / CLIGHT}') plt.xlabel("time [days]") plt.ylabel("Te [K]") @@ -79,7 +76,7 @@ def plot_Te_vs_velocity(modelpath, modeldata, estimators, readonly_mgi): associated_modeldata_rows = [ modeldata.loc[modeldata["inputcellid"] == assoc_cells[mgi][0]] for mgi in readonly_mgi ] - velocity = [row["vel_y_mid"].values[0] / CLIGHT for row in associated_modeldata_rows] + velocity = [row["vel_y_mid"].to_numpy()[0] / CLIGHT for row in associated_modeldata_rows] plt.plot(velocity, Te, label=f"{times[timestep]:.2f}", linestyle="-", marker="o") @@ -127,7 +124,14 @@ def make_2d_plot(grid, grid_Te, vmax, modelpath, xgrid, time): mesh = pv.StructuredGrid(x, y, z) mesh["Te [K]"] = grid_Te.ravel(order="F") - sargs = dict(height=0.75, vertical=True, position_x=0.02, position_y=0.1, title_font_size=22, label_font_size=25) + sargs = { + "height": 0.75, + "vertical": True, + "position_x": 0.02, + "position_y": 0.1, + "title_font_size": 22, + "label_font_size": 25, + } pv.set_plot_theme("document") # set white background p = pv.Plotter() diff --git a/artistools/estimators/plotestimators.py b/artistools/estimators/plotestimators.py old mode 100644 new mode 100755 index 480aeaaf2..2a45ae735 --- a/artistools/estimators/plotestimators.py +++ b/artistools/estimators/plotestimators.py @@ -4,7 +4,6 @@ Examples are temperatures, populations, heating/cooling rates. """ -# import math import argparse import math import multiprocessing @@ -18,8 +17,6 @@ import pandas as pd import artistools as at -import artistools.initial_composition -import artistools.nltepops colors_tab10 = list(plt.get_cmap("tab10")(np.linspace(0, 1.0, 10))) @@ -130,7 +127,7 @@ def plot_average_ionisation_excitation( elif seriestype == "averageexcitation": ax.set_ylabel("Average excitation [eV]") else: - raise ValueError() + raise ValueError arr_tdelta = at.get_timestep_times_float(modelpath, loc="delta") for paramvalue in params: @@ -188,8 +185,6 @@ def plot_levelpop( args=None, **plotkwargs, ): - import artistools.plottools - if seriestype == "levelpopulation_dn_on_dvel": ax.set_ylabel("dN/dV [{}km$^{{-1}}$ s]") ax.yaxis.set_major_formatter(at.plottools.ExponentLabelFormatter(ax.get_ylabel(), useMathText=True)) @@ -197,10 +192,10 @@ def plot_levelpop( ax.set_ylabel("X$_{{i}}$ [{}/cm3]") ax.yaxis.set_major_formatter(at.plottools.ExponentLabelFormatter(ax.get_ylabel(), useMathText=True)) else: - raise ValueError() + raise ValueError modeldata, _ = at.inputmodel.get_modeldata(modelpath) - modeldata.eval("modelcellvolume = cellmass_grams / (10 ** logrho)", inplace=True) + modeldata = modeldata.eval("modelcellvolume = cellmass_grams / (10 ** logrho)") adata = at.atomic.get_levels(modelpath) @@ -252,10 +247,11 @@ def plot_levelpop( if dfalldata is not None: elsym = at.get_elsymbol(atomic_number).lower() - if seriestype == "levelpopulation_dn_on_dvel": - colname = f"nlevel_on_dv_{elsym}_ionstage{ion_stage}_level{levelindex}" - else: - colname = f"nnlevel_{elsym}_ionstage{ion_stage}_level{levelindex}" + colname = ( + f"nlevel_on_dv_{elsym}_ionstage{ion_stage}_level{levelindex}" + if seriestype == "levelpopulation_dn_on_dvel" + else f"nnlevel_{elsym}_ionstage{ion_stage}_level{levelindex}" + ) dfalldata[colname] = ylist ylist.insert(0, ylist[0]) @@ -288,14 +284,13 @@ def plot_multi_ion_series( def get_iontuple(ionstr): if ionstr in at.get_elsymbolslist(): return (at.get_atomic_number(ionstr), "ALL") - elif " " in ionstr: + if " " in ionstr: return (at.get_atomic_number(ionstr.split(" ")[0]), at.decode_roman_numeral(ionstr.split(" ")[1])) - elif ionstr.rstrip("-0123456789") in at.get_elsymbolslist(): + if ionstr.rstrip("-0123456789") in at.get_elsymbolslist(): atomic_number = at.get_atomic_number(ionstr.rstrip("-0123456789")) return (atomic_number, ionstr) - else: - atomic_number = at.get_atomic_number(ionstr.split("_")[0]) - return (atomic_number, ionstr) + atomic_number = at.get_atomic_number(ionstr.split("_")[0]) + return (atomic_number, ionstr) # decoded into atomic number and parameter, e.g., [(26, 1), (26, 2), (26, 'ALL'), (26, 'Fe56')] iontuplelist = [get_iontuple(ionstr) for ionstr in ionlist] @@ -348,7 +343,7 @@ def get_iontuple(ionstr): elif args.ionpoptype == "totalpop": ax.set_ylabel(r"X$_{i}$/X$_{rm tot}$") else: - assert False + raise AssertionError else: ax.set_ylabel(at.estimators.get_dictlabelreplacements().get(seriestype, seriestype)) @@ -384,7 +379,7 @@ def get_iontuple(ionstr): totalpop = estimpop["total"] yvalue = nionpop / totalpop # Plot as fraction of total population else: - assert False + raise AssertionError except ZeroDivisionError: yvalue = 0.0 @@ -418,10 +413,11 @@ def get_iontuple(ionstr): yvalue = float("NaN") ylist.append(yvalue) - if hasattr(ion_stage, "lower") and ion_stage != "ALL": - plotlabel = ion_stage - else: - plotlabel = at.get_ionstring(atomic_number, ion_stage, spectral=False) + plotlabel = ( + ion_stage + if hasattr(ion_stage, "lower") and ion_stage != "ALL" + else at.get_ionstring(atomic_number, ion_stage, spectral=False) + ) color = get_elemcolor(atomic_number=atomic_number) @@ -544,12 +540,11 @@ def plot_series( ax.plot(xlist, ylist, linewidth=1.5, label=linelabel, color=dictcolors.get(variablename, None), **plotkwargs) -def get_xlist(xvariable, allnonemptymgilist, estimators, timestepslist, modelpath, args): +def get_xlist( + xvariable, allnonemptymgilist, estimators, timestepslist, modelpath, args +) -> tuple[list[float], list[float], list[float]]: if xvariable in ["cellid", "modelgridindex"]: - if args.xmax >= 0: - mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] - else: - mgilist_out = allnonemptymgilist + mgilist_out = [mgi for mgi in allnonemptymgilist if mgi <= args.xmax] if args.xmax >= 0 else allnonemptymgilist xlist = mgilist_out timestepslist_out = timestepslist elif xvariable == "timestep": @@ -573,7 +568,7 @@ def get_xlist(xvariable, allnonemptymgilist, estimators, timestepslist, modelpat if args.xmax > 0 and xvalue > args.xmax: break - xlist, mgilist_out, timestepslist_out = zip(*list(sorted(zip(xlist, mgilist_out, timestepslist_out)))) + xlist, mgilist_out, timestepslist_out = zip(*sorted(zip(xlist, mgilist_out, timestepslist_out))) assert len(xlist) == len(mgilist_out) == len(timestepslist_out) @@ -588,12 +583,12 @@ def plot_subplot( assert len(xlist) - 1 == len(mgilist) == len(timestepslist) showlegend = False - ylabel = "UNDEFINED" + ylabel = None sameylabel = True for variablename in plotitems: if not isinstance(variablename, str): pass - elif ylabel == "UNDEFINED": + elif ylabel is None: ylabel = get_ylabel(variablename) elif ylabel != get_ylabel(variablename): sameylabel = False @@ -601,7 +596,7 @@ def plot_subplot( for plotitem in plotitems: if isinstance(plotitem, str): - showlegend = len(plotitems) > 1 or len(variablename) > 20 + showlegend = len(plotitems) > 1 or len(plotitem) > 20 plot_series( ax, xlist, @@ -721,10 +716,7 @@ def make_plot( dfalldata.index.name = "modelgridindex" dfalldata[xvariable] = xlist - if xvariable.startswith("velocity"): - xlist = np.insert(xlist, 0, 0.0) - else: - xlist = np.insert(xlist, 0, xlist[0]) + xlist = np.insert(xlist, 0, 0.0) if xvariable.startswith("velocity") else np.insert(xlist, 0, xlist[0]) xmin = args.xmin if args.xmin >= 0 else min(xlist) xmax = args.xmax if args.xmax > 0 else max(xlist) @@ -775,7 +767,7 @@ def make_plot( # plt.suptitle(figure_title, fontsize=11, verticalalignment='top') if args.write_data: - dfalldata.sort_index(inplace=True) + dfalldata = dfalldata.sort_index() dataoutfilename = Path(outfilename).with_suffix(".txt") dfalldata.to_csv(dataoutfilename) print(f"Saved {dataoutfilename}") @@ -915,12 +907,12 @@ def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument( "-filtersavgol", nargs=2, - help="Savitzky–Golay filter. Specify the window_length and polyorder.e.g. -filtersavgol 5 3", + help="Savitzky-Golay filter. Specify the window_length and polyorder.e.g. -filtersavgol 5 3", ) parser.add_argument("--notitle", action="store_true", help="Suppress the top title from the plot") - parser.add_argument("-plotlist", type=list, default=[], help="Plot list (when calling from Python only)") # type: ignore + parser.add_argument("-plotlist", type=list, default=[], help="Plot list (when calling from Python only)") # type: ignore[arg-type] parser.add_argument( "-ionpoptype", @@ -993,7 +985,7 @@ def main(args=None, argsraw=None, **kwargs): ) for ts in reversed(timesteps_included): - tswithdata = [ts for (ts, mgi) in estimators.keys()] + tswithdata = [ts for (ts, mgi) in estimators] for ts in timesteps_included: if ts not in tswithdata: timesteps_included.remove(ts) @@ -1003,10 +995,10 @@ def main(args=None, argsraw=None, **kwargs): print("No timesteps with data are included") return - if args.plotlist: - plotlist = args.plotlist - else: - plotlist = [ + plotlist = ( + args.plotlist + if args.plotlist + else [ # [['initabundances', ['Fe', 'Ni_stable', 'Ni_56']]], # ['heating_dep', 'heating_coll', 'heating_bf', 'heating_ff', # ['_yscale', 'linear']], @@ -1039,6 +1031,7 @@ def main(args=None, argsraw=None, **kwargs): # [['Alpha_R / RRC_LTE_Nahar', ['Fe II', 'Fe III', 'Fe IV', 'Fe V', 'Ni III']]], # [['gamma_NT', ['Fe I', 'Fe II', 'Fe III', 'Fe IV', 'Fe V', 'Ni II']]], ] + ) if args.recombrates: plot_recombrates(modelpath, estimators, 26, [2, 3, 4, 5]) diff --git a/artistools/gsinetwork.py b/artistools/gsinetwork.py index f3db75170..0ea4b046b 100755 --- a/artistools/gsinetwork.py +++ b/artistools/gsinetwork.py @@ -17,11 +17,6 @@ import artistools as at -# import io -# import math - -# import artistools.estimators - def plot_qdot( modelpath: Path, @@ -204,7 +199,7 @@ def plot_qdot( # fig.suptitle(f'{modelname}', fontsize=10) at.plottools.autoscale(axis, margin=0.0) - plt.savefig(pdfoutpath, format="pdf") + fig.savefig(pdfoutpath, format="pdf") print(f"Saved {pdfoutpath}") @@ -310,7 +305,7 @@ def plot_cell_abund_evolution( fig.suptitle(f"{modelname} cell {mgi}", y=0.995, fontsize=10) at.plottools.autoscale(axis, margin=0.05) - plt.savefig(pdfoutpath, format="pdf") + fig.savefig(pdfoutpath, format="pdf") print(f"Saved {pdfoutpath}") @@ -415,7 +410,7 @@ def plot_qdot_abund_modelcells(modelpath: Path, mgiplotlist: Sequence[int], arr_ griddatafolder: Path = Path("SFHo_snapshot") mergermodelfolder: Path = Path("SFHo_short") trajfolder: Path = Path("SFHo") - with at.zopen(modelpath / "model.txt", "rt") as fmodel: + with at.zopen(modelpath / "model.txt") as fmodel: while True: line = fmodel.readline() if line.startswith("#"): @@ -442,7 +437,7 @@ def plot_qdot_abund_modelcells(modelpath: Path, mgiplotlist: Sequence[int], arr_ dfmodel, t_model_init_days, vmax_cmps = at.inputmodel.get_modeldata_tuple(modelpath) if "logrho" not in dfmodel.columns: - dfmodel.eval("logrho = log10(rho)", inplace=True) + dfmodel = dfmodel.eval("logrho = log10(rho)") model_mass_grams = dfmodel.cellmass_grams.sum() npts_model = len(dfmodel) @@ -452,12 +447,12 @@ def plot_qdot_abund_modelcells(modelpath: Path, mgiplotlist: Sequence[int], arr_ # WARNING sketchy inference! propcellcount = math.ceil(max(mgi_of_propcells.keys()) ** (1 / 3.0)) ** 3 xmax_tmodel = vmax_cmps * t_model_init_days * 86400 - wid_init = at.misc.get_wid_init_at_tmodel(modelpath, propcellcount, t_model_init_days, xmax_tmodel) + wid_init = at.get_wid_init_at_tmodel(modelpath, propcellcount, t_model_init_days, xmax_tmodel) dfmodel["n_assoc_cells"] = [len(assoc_cells.get(inputcellid - 1, [])) for inputcellid in dfmodel["inputcellid"]] # for spherical models, ARTIS mapping to a cubic grid introduces some errors in the cell volumes - dfmodel.eval("cellmass_grams_mapped = 10 ** logrho * @wid_init ** 3 * n_assoc_cells", inplace=True) - for strnuc, a in zip(arr_strnuc, arr_a): + dfmodel = dfmodel.eval("cellmass_grams_mapped = 10 ** logrho * @wid_init ** 3 * n_assoc_cells") + for strnuc in arr_strnuc: corr = ( dfmodel.eval(f"X_{strnuc} * cellmass_grams_mapped").sum() / dfmodel.eval(f"X_{strnuc} * cellmass_grams").sum() @@ -565,7 +560,7 @@ def plot_qdot_abund_modelcells(modelpath: Path, mgiplotlist: Sequence[int], arr_ arr_time_gsi_days = list(arr_time_gsi_s / 86400) dfpartcontrib = at.inputmodel.rprocess_from_trajectory.get_gridparticlecontributions(modelpath) - dfpartcontrib.query("cellindex <= @npts_model and frac_of_cellmass > 0", inplace=True) + dfpartcontrib = dfpartcontrib.query("cellindex <= @npts_model and frac_of_cellmass > 0") list_particleids_getabund = dfpartcontrib.query("(cellindex - 1) in @mgiplotlist").particleid.unique() fworkerwithabund = partial(get_particledata, arr_time_gsi_s_incpremerger, arr_strnuc, traj_root) @@ -595,9 +590,7 @@ def plot_qdot_abund_modelcells(modelpath: Path, mgiplotlist: Sequence[int], arr_ else: list_particledata_noabund = [fworkernoabund(particleid) for particleid in list_particleids_noabund] - allparticledata = { - particleid: data for particleid, data in (list_particledata_withabund + list_particledata_noabund) - } + allparticledata = dict(list_particledata_withabund + list_particledata_noabund) plot_qdot( modelpath, diff --git a/artistools/hesma_scripts.py b/artistools/hesma_scripts.py index b186032b7..7008f11f2 100644 --- a/artistools/hesma_scripts.py +++ b/artistools/hesma_scripts.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 import os from pathlib import Path @@ -8,9 +7,6 @@ import artistools as at -# import artistools.spectra -# import artistools.lightcurve.writebollightcurvedata - def plot_hesma_spectrum(timeavg, axes): hesma_file = Path("/Users/ccollins/Downloads/hesma_files/M2a/hesma_specseq.dat") @@ -48,13 +44,13 @@ def plothesmaresspec(fig, ax): # chunk = specdata.iloc[index_to_split[i]:, :] # res_specdata.append(chunk) - res_specdata = at.gather_res_data(specdata) + res_specdata = at.split_dataframe_dirbins(specdata) column_names = res_specdata[0].iloc[0] column_names[0] = "lambda" print(column_names) - for i, res_spec in enumerate(res_specdata): + for i, _res_spec in enumerate(res_specdata): res_specdata[i] = res_specdata[i].rename(columns=column_names).drop(res_specdata[i].index[0]) ax.plot(res_specdata[0]["lambda"], res_specdata[0][11.7935] * (1e-5) ** 2, label="hesma 0") @@ -81,9 +77,9 @@ def make_hesma_vspecfiles(modelpath, outpath=None): vspecdata_all = at.spectra.get_specpol_data(angle=angle, modelpath=modelpath) vspecdata = vspecdata_all["I"] - timearray = vspecdata.columns.values[1:] - vspecdata.sort_values(by="nu", ascending=False, inplace=True) - vspecdata.eval("lambda_angstroms = 2.99792458e+18 / nu", inplace=True) + timearray = vspecdata.columns.to_numpy()[1:] + vspecdata = vspecdata.sort_values(by="nu", ascending=False) + vspecdata = vspecdata.eval("lambda_angstroms = 2.99792458e+18 / nu") for time in timearray: vspecdata[time] = vspecdata[time] * vspecdata["nu"] / vspecdata["lambda_angstroms"] vspecdata[time] = vspecdata[time] * (1e5) ** 2 # Scale to 10 pc (1 Mpc/10 pc) ** 2 @@ -106,7 +102,9 @@ def make_hesma_vspecfiles(modelpath, outpath=None): f.write( f"# File contains spectra at observer angles {angle_names} for Model {modelname}.\n# A header line" " containing spectra time is repeated at the beginning of each observer angle. Column 0 gives wavelength." - " \n# Spectra are at a distance of 10 pc." + "\n" + content + " \n# Spectra are at a distance of 10 pc." + "\n" + + content ) @@ -115,7 +113,7 @@ def make_hesma_bol_lightcurve(modelpath, outpath, timemin, timemax): lightcurvedataframe = at.lightcurve.writebollightcurvedata.get_bol_lc_from_lightcurveout(modelpath) print(lightcurvedataframe) - lightcurvedataframe.query("time > @timemin and time < @timemax", inplace=True) + lightcurvedataframe = lightcurvedataframe.query("time > @timemin and time < @timemax") modelname = at.get_model_name(modelpath) outfilename = f"doubledet_2021_{modelname}.dat" @@ -144,9 +142,7 @@ def make_hesma_peakmag_dm15_dm40(band, pathtofiles, modelname, outpath, dm40=Fal ) angles = np.arange(0, 100) - angle_definition = at.lightcurve.viewingangleanalysis.calculate_costheta_phi_for_viewing_angles( - angles, modelpath=None - ) + angle_definition = at.get_dirbin_labels(angles, modelpath=None) outdata = {} outdata["peakmag"] = dm15data["peakmag"] # dm15 peak mag probably more accurate - shorter time window diff --git a/artistools/initial_composition.py b/artistools/initial_composition.py old mode 100644 new mode 100755 index b76a3e37f..d8a070079 --- a/artistools/initial_composition.py +++ b/artistools/initial_composition.py @@ -1,32 +1,32 @@ #!/usr/bin/env python3 +# PYTHON_ARGCOMPLETE_OK import argparse import math import os from pathlib import Path +from typing import Any +from typing import Literal +from typing import Optional -import matplotlib +import argcomplete +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np +import pandas as pd from astropy import units as u import artistools as at -import artistools.inputmodel.opacityinputfile - -# import pandas as pd - -# import artistools.inputmodel -# from mpl_toolkits.mplot3d import Axes3D def plot_2d_initial_abundances(modelpath, args): - model = at.inputmodel.get_2d_modeldata(modelpath[0]) - abundances = at.inputmodel.get_initialabundances(modelpath[0]) + model = at.inputmodel.get_2d_modeldata(modelpath) + abundances = at.inputmodel.get_initelemabundances(modelpath) - abundances["inputcellid"] = abundances["inputcellid"].apply(lambda x: float(x)) + abundances["inputcellid"] = abundances["inputcellid"].apply(float) merge_dfs = model.merge(abundances, how="inner", on="inputcellid") - with open(os.path.join(modelpath[0], "model.txt"), "r") as fmodelin: + with open(os.path.join(modelpath, "model.txt")) as fmodelin: fmodelin.readline() # npts r, npts z t_model = float(fmodelin.readline()) # days vmax = float(fmodelin.readline()) # v_max in [cm/s] @@ -34,64 +34,56 @@ def plot_2d_initial_abundances(modelpath, args): r = merge_dfs["cellpos_mid[r]"] / t_model * (u.cm / u.day).to("km/s") / 10**3 z = merge_dfs["cellpos_mid[z]"] / t_model * (u.cm / u.day).to("km/s") / 10**3 - ion = f"X_{args.ion}" + colname = f"X_{args.plotvars}" font = {"weight": "bold", "size": 18} f = plt.figure(figsize=(4, 5)) ax = f.add_subplot(111) - im = ax.scatter(r, z, c=merge_dfs[ion], marker="8") + im = ax.scatter(r, z, c=merge_dfs[colname], marker="8") f.colorbar(im) plt.xlabel(r"v$_x$ in 10$^3$ km/s", fontsize="x-large") # , fontweight='bold') plt.ylabel(r"v$_z$ in 10$^3$ km/s", fontsize="x-large") # , fontweight='bold') - plt.text(20, 25, args.ion, color="white", fontweight="bold", fontsize="x-large") + plt.text(20, 25, args.plotvars, color="white", fontweight="bold", fontsize="x-large") plt.tight_layout() # ax.labelsize: 'large' # plt.title(f'At {sliceaxis} = {sliceposition}') - outfilename = f"plotcomposition{args.ion}.pdf" - plt.savefig(Path(modelpath[0]) / outfilename, format="pdf") + outfilename = f"plotcomposition{args.plotvars}.pdf" + plt.savefig(Path(modelpath) / outfilename, format="pdf") print(f"Saved {outfilename}") -def get_merged_model_abundances(modelpath): - # t_model is in days and vmax is in cm/s - model, t_model_days, vmax = at.inputmodel.get_modeldata_tuple(modelpath[0], dimensions=3) - - targetmodeltime_days = None - if targetmodeltime_days is not None: - print( - f"Scaling modeldata to {targetmodeltime_days} days. \nWARNING: abundances not scaled for radioactive decays" - ) - import artistools.inputmodel.modelfromhydro - - artistools.inputmodel.modelfromhydro.scale_model_to_time(targetmodeltime_days, t_model_days, model) - t_model_days = targetmodeltime_days - - abundances = at.inputmodel.get_initialabundances(modelpath[0]) - - abundances["inputcellid"] = abundances["inputcellid"].apply(lambda x: float(x)) - - merge_dfs = model.merge(abundances, how="inner", on="inputcellid") - return merge_dfs, t_model_days - - -def get_2D_slice_through_3d_model(merge_dfs, sliceaxis, sliceindex=None): +def get_2D_slice_through_3d_model( + dfmodel: pd.DataFrame, + sliceaxis: Literal["x", "y", "z"], + modelmeta: Optional[dict[str, Any]] = None, + plotaxis1: Optional[Literal["x", "y", "z"]] = None, + plotaxis2: Optional[Literal["x", "y", "z"]] = None, + sliceindex: Optional[int] = None, +) -> pd.DataFrame: if not sliceindex: # get midpoint - sliceposition = merge_dfs.iloc[(merge_dfs["pos_x_min"]).abs().argsort()][:1]["pos_x_min"].item() + sliceposition: float = dfmodel.iloc[(dfmodel["pos_x_min"]).abs().argsort()][:1]["pos_x_min"].item() # Choose position to slice. This gets minimum absolute value as the closest to 0 else: cell_boundaries = [] - [cell_boundaries.append(x) for x in merge_dfs[f"pos_{sliceaxis}_min"] if x not in cell_boundaries] + for x in dfmodel[f"pos_{sliceaxis}_min"]: + if x not in cell_boundaries: + cell_boundaries.append(x) sliceposition = cell_boundaries[sliceindex] - slicedf = merge_dfs.loc[merge_dfs[f"pos_{sliceaxis}_min"] == sliceposition] + slicedf = dfmodel.loc[dfmodel[f"pos_{sliceaxis}_min"] == sliceposition] + + if modelmeta is not None and plotaxis1 is not None and plotaxis2 is not None: + assert slicedf.shape[0] == modelmeta[f"ncoordgrid{plotaxis1}"] * modelmeta[f"ncoordgrid{plotaxis2}"] + return slicedf -def plot_abundances_ion(ax, plotvals, ion, plotaxis1, plotaxis2, t_model, args): - colorscale = plotvals[ion] +def plot_slice_modelcol(ax, plotvals, modelmeta, colname, plotaxis1, plotaxis2, t_model_d, args): + print(colname) + colorscale = plotvals[colname] * plotvals["rho"] if colname.startswith("X_") else plotvals[colname] if args.hideemptycells: # Don't plot empty cells: @@ -100,57 +92,106 @@ def plot_abundances_ion(ax, plotvals, ion, plotaxis1, plotaxis2, t_model, args): if args.logcolorscale: # logscale for colormap colorscale = np.log10(colorscale) + colorscale = colorscale.to_numpy() normalise_between_0_and_1 = False if normalise_between_0_and_1: - norm = matplotlib.colors.Normalize(vmin=0, vmax=1) - scaledmap = matplotlib.cm.ScalarMappable(cmap="viridis", norm=norm) + norm = mpl.colors.Normalize(vmin=0, vmax=1) + scaledmap = mpl.cm.ScalarMappable(cmap="viridis", norm=norm) scaledmap.set_array([]) colorscale = scaledmap.to_rgba(colorscale) # colorscale fixed between 0 and 1 else: scaledmap = None - x = plotvals[f"pos_{plotaxis1}_min"] / t_model * (u.cm / u.day).to("km/s") / 10**3 - y = plotvals[f"pos_{plotaxis2}_min"] / t_model * (u.cm / u.day).to("km/s") / 10**3 + arr_x = plotvals[f"pos_{plotaxis1}_min"] / t_model_d / 86400 / 2.99792458e10 + arr_y = plotvals[f"pos_{plotaxis2}_min"] / t_model_d / 86400 / 2.99792458e10 # x = plotvals[f'pos_{plotaxis1}'] / t_model * (u.cm/u.day).to('m/s') / 2.99792458e+8 # y = plotvals[f'pos_{plotaxis2}'] / t_model * (u.cm/u.day).to('m/s') / 2.99792458e+8 - im = ax.scatter(x, y, c=colorscale, marker="8", rasterized=True) # cmap=plt.get_cmap('PuOr') + # im = ax.scatter(x, y, c=colorscale, marker="s", s=30, rasterized=False) # cmap=plt.get_cmap('PuOr') + ncoordgrid1 = modelmeta[f"ncoordgrid{plotaxis1}"] + ncoordgrid2 = modelmeta[f"ncoordgrid{plotaxis2}"] + grid = np.zeros((ncoordgrid1, ncoordgrid2)) + + for i in range(0, ncoordgrid1): + for j in range(0, ncoordgrid2): + grid[j, i] = colorscale[j * ncoordgrid1 + i] + + im = ax.imshow( + grid, + cmap="viridis", + interpolation="nearest", + extent=(arr_x.min(), arr_x.max(), arr_y.min(), arr_y.max()), + origin="lower", + # vmin=0.0, + # vmax=1.0, + # vmax=-9.5, + # vmin=-11, + # vmin=1e-11, + ) + + plot_vmax = 0.2 + ax.set_ylim(bottom=-plot_vmax, top=plot_vmax) + ax.set_xlim(left=-plot_vmax, right=plot_vmax) + if "_" in colname: + ax.annotate( + colname.split("_")[1], + color="white", + xy=(0.9, 0.9), + xycoords="axes fraction", + horizontalalignment="right", + verticalalignment="top", + # fontsize=10, + ) - ymin, ymax = ax.get_ylim() - xmin, xmax = ax.get_xlim() - if "_" in ion: - ax.text(xmax * 0.6, ymax * 0.7, ion.split("_")[1], color="k", fontweight="bold") return im, scaledmap -def plot_3d_initial_abundances(modelpath, args=None): +def plot_3d_initial_abundances(modelpath, args=None) -> None: font = { # 'weight': 'bold', "size": 18 } - matplotlib.rc("font", **font) + mpl.rc("font", **font) - merge_dfs, t_model = get_merged_model_abundances(modelpath) - # merge_dfs = plot_most_abundant(modelpath, args) + dfmodel, modelmeta = at.get_modeldata( + modelpath, skipnuclidemassfraccolumns=True, get_elemabundances=True, dtype_backend="pyarrow" + ) + + targetmodeltime_days = None + if targetmodeltime_days is not None: + print( + f"Scaling modeldata to {targetmodeltime_days} days. \nWARNING: abundances not scaled for radioactive decays" + ) + + dfmodel, modelmeta = at.inputmodel.scale_model_to_time( + targetmodeltime_days=targetmodeltime_days, modelmeta=modelmeta, dfmodel=dfmodel + ) + + sliceaxis: Literal["x", "y", "z"] = "z" - plotaxis1 = "y" - plotaxis2 = "z" - sliceaxis = "x" + axes: list[Literal["x", "y", "z"]] = ["x", "y", "z"] + plotaxis1: Literal["x", "y", "z"] = [ax for ax in axes if ax != sliceaxis][0] + plotaxis2: Literal["x", "y", "z"] = [ax for ax in axes if ax not in [sliceaxis, plotaxis1]][0] - plotvals = get_2D_slice_through_3d_model(merge_dfs, sliceaxis) + df2dslice = get_2D_slice_through_3d_model( + dfmodel=dfmodel, modelmeta=modelmeta, sliceaxis=sliceaxis, plotaxis1=plotaxis1, plotaxis2=plotaxis2 + ) subplots = False - if len(args.ion) > 1: + if len(args.plotvars) > 1: subplots = True if not subplots: - fig = plt.figure(figsize=(8, 7)) + fig = plt.figure( + figsize=(8, 7), + tight_layout={"pad": 0.4, "w_pad": 0.0, "h_pad": 0.0}, + ) ax = plt.subplot(111, aspect="equal") else: rows = 1 - cols = len(args.ion) + cols = len(args.plotvars) fig, axes = plt.subplots( nrows=rows, @@ -158,24 +199,25 @@ def plot_3d_initial_abundances(modelpath, args=None): sharex=True, sharey=True, figsize=(at.get_config()["figwidth"] * cols, at.get_config()["figwidth"] * 1.4), - tight_layout={"pad": 5.0, "w_pad": 0.0, "h_pad": 0.0}, + # tight_layout={"pad": 5.0, "w_pad": 0.0, "h_pad": 0.0}, ) for ax in axes: ax.set(aspect="equal") - for index, ion in enumerate(args.ion): - ion = f"X_{ion}" - if args.rho: - ion = "rho" + for index, plotvar in enumerate(args.plotvars): + colname = plotvar if plotvar in df2dslice.columns else f"X_{plotvar}" + if subplots: ax = axes[index] - im, scaledmap = plot_abundances_ion(ax, plotvals, ion, plotaxis1, plotaxis2, t_model, args) + im, scaledmap = plot_slice_modelcol( + ax, df2dslice, modelmeta, colname, plotaxis1, plotaxis2, modelmeta["t_model_init_days"], args + ) - xlabel = rf"v$_{plotaxis1}$ in 10$^3$ km/s" - ylabel = rf"v$_{plotaxis2}$ in 10$^3$ km/s" + xlabel = rf"v$_{plotaxis1}$ [$c$]" + ylabel = rf"v$_{plotaxis2}$ [$c$]" if not subplots: - cbar = plt.colorbar(im) + cbar = fig.colorbar(im) plt.xlabel(xlabel, fontsize="x-large") # , fontweight='bold') plt.ylabel(ylabel, fontsize="x-large") # , fontweight='bold') else: @@ -183,19 +225,20 @@ def plot_3d_initial_abundances(modelpath, args=None): fig.text(0.5, 0.15, xlabel, ha="center", va="center") fig.text(0.05, 0.5, ylabel, ha="center", va="center", rotation="vertical") - # cbar.set_label(label=ion, size='x-large') #, fontweight='bold') - # cbar.ax.set_title(f'{args.ion}', size='small') + # cbar.set_label(label="test", size="x-large") # , fontweight='bold') + if "cellYe" not in args.plotvars and "tracercount" not in args.plotvars: + if args.logcolorscale: + cbar.ax.set_title(r"log10($\rho$) [g/cm3]", size="small") + else: + cbar.ax.set_title(r"$\rho$ [g/cm3]", size="small") # cbar.ax.tick_params(labelsize='x-large') # plt.tight_layout() # ax.labelsize: 'large' # plt.title(f'At {sliceaxis} = {sliceposition}') - # if args.outputfile: - # outfilename = args.outputfile - # else: - outfilename = f"plotcomposition{ion}.pdf" - plt.savefig(Path(modelpath[0]) / outfilename, format="pdf") + outfilename = args.outputfile if args.outputfile else f"plotcomposition_{','.join(args.plotvars)}.pdf" + plt.savefig(Path(modelpath) / outfilename, format="pdf") print(f"Saved {outfilename}") @@ -203,7 +246,7 @@ def plot_3d_initial_abundances(modelpath, args=None): def get_model_abundances_Msun_1D(modelpath): filename = modelpath / "model.txt" modeldata, t_model_init_days, _ = at.inputmodel.get_modeldata_tuple(filename) - abundancedata = at.inputmodel.get_initialabundances(modelpath) + abundancedata = at.inputmodel.get_initelemabundances(modelpath) t_model_init_seconds = t_model_init_days * 24 * 60 * 60 @@ -222,7 +265,7 @@ def get_model_abundances_Msun_1D(modelpath): merge_dfs = modeldata.merge(abundancedata, how="inner", on="inputcellid") print("Total mass (Msun):") - for key in merge_dfs.keys(): + for key in merge_dfs: if "X_" in key: merge_dfs[f"mass_{key}"] = merge_dfs[key] * merge_dfs["mass_shell"] * u.g.to("solMass") # get mass of element in each cell @@ -232,11 +275,11 @@ def get_model_abundances_Msun_1D(modelpath): def plot_most_abundant(modelpath, args): - model, _ = at.inputmodel.get_modeldata(modelpath[0], dimensions=3) - abundances = at.inputmodel.get_initialabundances(modelpath[0]) + model, _ = at.inputmodel.get_modeldata(modelpath[0]) + abundances = at.inputmodel.get_initelemabundances(modelpath[0]) merge_dfs = model.merge(abundances, how="inner", on="inputcellid") - elements = [x for x in merge_dfs.keys() if "X_" in x] + elements = [x for x in merge_dfs if "X_" in x] merge_dfs["max"] = merge_dfs[elements].idxmax(axis=1) @@ -251,10 +294,10 @@ def make_3d_plot(modelpath, args): pv.set_plot_theme("document") # set white background - model, t_model, vmax = at.inputmodel.get_modeldata_tuple(modelpath, dimensions=3, get_elemabundances=False) - abundances = at.inputmodel.get_initialabundances(modelpath) + model, t_model, vmax = at.inputmodel.get_modeldata_tuple(modelpath, get_elemabundances=False) + abundances = at.inputmodel.get_initelemabundances(modelpath) - abundances["inputcellid"] = abundances["inputcellid"].apply(lambda x: float(x)) + abundances["inputcellid"] = abundances["inputcellid"].apply(float) merge_dfs = model.merge(abundances, how="inner", on="inputcellid") model = merge_dfs @@ -266,8 +309,8 @@ def make_3d_plot(modelpath, args): model["opacity"] = at.inputmodel.opacityinputfile.get_opacity_from_file(modelpath) coloursurfaceby = "opacity" else: - print(f"Colours set by X_{args.ion}") - coloursurfaceby = f"X_{args.ion}" + print(f"Colours set by X_{args.plotvars}") + coloursurfaceby = f"X_{args.plotvars}" # generate grid from data grid = round(len(model["rho"]) ** (1.0 / 3.0)) @@ -309,28 +352,30 @@ def make_3d_plot(modelpath, args): def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument( "-modelpath", - default=[], - nargs="*", - action=at.AppendPath, - help="Path(s) to ARTIS folder (may include wildcards such as * and **)", + type=Path, + default=Path(), + help="Path to ARTIS folder", ) + parser.add_argument("-o", action="store", dest="outputfile", type=Path, default=None, help="Filename for PDF file") + parser.add_argument( - "-o", action="store", dest="outputfile", type=Path, default=Path(), help="Filename for PDF file" + "plotvars", + type=str, + default=["rho"], + nargs="+", + help=( + "Element symbols (Fe, Ni, Sr) for mass fraction or other model columns (rho, tracercount) to plot. Default" + " is rho" + ), ) - parser.add_argument("-ion", type=str, default=["Fe"], nargs="+", help="Choose ion to plot. Default is Fe") - - parser.add_argument("--rho", action="store_true", help="Plot rho instead of ion") - parser.add_argument("--logcolorscale", action="store_true", help="Use log scale for colour map") parser.add_argument("--hideemptycells", action="store_true", help="Don't plot empty cells") parser.add_argument("--opacity", action="store_true", help="Plot opacity from opacity.txt (if available for model)") - parser.add_argument("-modeldim", type=int, default=None, help="Choose how many dimensions. 3 for 3D, 2 for 2D") - parser.add_argument("--plot3d", action="store_true", help="Make 3D plot") parser.add_argument("-surfaces3d", type=float, nargs="+", help="define positions of surfaces for 3D plots") @@ -343,26 +388,22 @@ def main(args=None, argsraw=None, **kwargs): ) addargs(parser) parser.set_defaults(**kwargs) + argcomplete.autocomplete(parser) args = parser.parse_args(argsraw) if not args.modelpath: args.modelpath = ["."] - args.modelpath = at.flatten_list(args.modelpath) - if args.plot3d: - make_3d_plot(Path(args.modelpath[0]), args) + make_3d_plot(Path(args.modelpath), args) return - if not args.modeldim: - inputparams = at.get_inputparams(args.modelpath[0]) - else: - inputparams = {"n_dimensions": args.modeldim} + _, modelmeta = at.get_modeldata(getheadersonly=True, printwarningsonly=True) - if inputparams["n_dimensions"] == 2: + if modelmeta["dimensions"] == 2: plot_2d_initial_abundances(args.modelpath, args) - if inputparams["n_dimensions"] == 3: + elif modelmeta["dimensions"] == 3: plot_3d_initial_abundances(args.modelpath, args) diff --git a/artistools/inputmodel/1dslicefrom3d.py b/artistools/inputmodel/1dslicefrom3d.py index a9d9915d9..5697955c0 100755 --- a/artistools/inputmodel/1dslicefrom3d.py +++ b/artistools/inputmodel/1dslicefrom3d.py @@ -35,13 +35,12 @@ def main(args=None, argsraw=None, **kwargs): if not os.path.exists(args.outputfolder): os.makedirs(args.outputfolder) - else: - if os.path.exists(os.path.join(args.outputfolder, "model.txt")): - print("ABORT: model.txt already exists") - sys.exit() - elif os.path.exists(os.path.join(args.outputfolder, "abundances.txt")): - print("ABORT: abundances.txt already exists") - sys.exit() + elif os.path.exists(os.path.join(args.outputfolder, "model.txt")): + print("ABORT: model.txt already exists") + sys.exit() + elif os.path.exists(os.path.join(args.outputfolder, "abundances.txt")): + print("ABORT: abundances.txt already exists") + sys.exit() dict3dcellidto1dcellid, xlist, ylists = slice_3dmodel(args.inputfolder, args.outputfolder, args.chosenaxis) @@ -57,7 +56,7 @@ def slice_3dmodel(inputfolder, outputfolder, chosenaxis): listout = [] dict3dcellidto1dcellid = {} outcellid = 0 - with open(os.path.join(inputfolder, "model.txt"), "r") as fmodelin: + with open(os.path.join(inputfolder, "model.txt")) as fmodelin: fmodelin.readline() # npts_model3d t_model = fmodelin.readline() # days fmodelin.readline() # v_max in [cm/s] @@ -83,11 +82,11 @@ def slice_3dmodel(inputfolder, outputfolder, chosenaxis): print("Wrong line size") sys.exit() - if cell["pos_x_min"] != "0.0000000" and (chosenaxis != "x" or float(cell["pos_x_min"]) < 0.0): - pass - elif cell["pos_y_min"] != "0.0000000" and (chosenaxis != "y" or float(cell["pos_y_min"]) < 0.0): - pass - elif cell["pos_z_min"] != "0.0000000" and (chosenaxis != "z" or float(cell["pos_z_min"]) < 0.0): + if ( + (cell["pos_x_min"] != "0.0000000" and (chosenaxis != "x" or float(cell["pos_x_min"]) < 0.0)) + or (cell["pos_y_min"] != "0.0000000" and (chosenaxis != "y" or float(cell["pos_y_min"]) < 0.0)) + or (cell["pos_z_min"] != "0.0000000" and (chosenaxis != "z" or float(cell["pos_z_min"]) < 0.0)) + ): pass else: outcellid += 1 @@ -108,7 +107,7 @@ def slice_3dmodel(inputfolder, outputfolder, chosenaxis): def slice_abundance_file(inputfolder, outputfolder, dict3dcellidto1dcellid): with ( - open(os.path.join(inputfolder, "abundances.txt"), "r") as fabundancesin, + open(os.path.join(inputfolder, "abundances.txt")) as fabundancesin, open(os.path.join(outputfolder, "abundances.txt"), "w") as fabundancesout, ): currentblock = [] diff --git a/artistools/inputmodel/__init__.py b/artistools/inputmodel/__init__.py index 47d0b3318..4566a7650 100644 --- a/artistools/inputmodel/__init__.py +++ b/artistools/inputmodel/__init__.py @@ -1,19 +1,23 @@ import artistools.inputmodel.botyanski2017 import artistools.inputmodel.describeinputmodel +import artistools.inputmodel.energyinputfiles import artistools.inputmodel.makeartismodel +import artistools.inputmodel.modelfromhydro +import artistools.inputmodel.opacityinputfile import artistools.inputmodel.rprocess_from_trajectory -from artistools.inputmodel.inputmodel_misc import add_derived_cols_to_modeldata -from artistools.inputmodel.inputmodel_misc import get_2d_modeldata -from artistools.inputmodel.inputmodel_misc import get_3d_model_data_merged_model_and_abundances_minimal -from artistools.inputmodel.inputmodel_misc import get_3d_modeldata_minimal -from artistools.inputmodel.inputmodel_misc import get_cell_angle -from artistools.inputmodel.inputmodel_misc import get_dfmodel_dimensions -from artistools.inputmodel.inputmodel_misc import get_initialabundances -from artistools.inputmodel.inputmodel_misc import get_mean_cell_properties_of_angle_bin -from artistools.inputmodel.inputmodel_misc import get_mgi_of_velocity_kms -from artistools.inputmodel.inputmodel_misc import get_modeldata -from artistools.inputmodel.inputmodel_misc import get_modeldata_tuple -from artistools.inputmodel.inputmodel_misc import save_empty_abundance_file -from artistools.inputmodel.inputmodel_misc import save_initialabundances -from artistools.inputmodel.inputmodel_misc import save_modeldata -from artistools.inputmodel.inputmodel_misc import sphericalaverage +from .inputmodel_misc import add_derived_cols_to_modeldata +from .inputmodel_misc import get_2d_modeldata +from .inputmodel_misc import get_3d_model_data_merged_model_and_abundances_minimal +from .inputmodel_misc import get_3d_modeldata_minimal +from .inputmodel_misc import get_cell_angle +from .inputmodel_misc import get_dfmodel_dimensions +from .inputmodel_misc import get_initelemabundances +from .inputmodel_misc import get_mean_cell_properties_of_angle_bin +from .inputmodel_misc import get_mgi_of_velocity_kms +from .inputmodel_misc import get_modeldata +from .inputmodel_misc import get_modeldata_tuple +from .inputmodel_misc import save_empty_abundance_file +from .inputmodel_misc import save_initelemabundances +from .inputmodel_misc import save_modeldata +from .inputmodel_misc import scale_model_to_time +from .inputmodel_misc import sphericalaverage diff --git a/artistools/inputmodel/botyanski2017.py b/artistools/inputmodel/botyanski2017.py index 926a3fc7b..3a1a80631 100755 --- a/artistools/inputmodel/botyanski2017.py +++ b/artistools/inputmodel/botyanski2017.py @@ -8,7 +8,6 @@ from astropy import units as u import artistools as at -import artistools.inputmodel def min_dist(listin, number): @@ -71,7 +70,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: fixed_points = [v_transition, v_ni56] regular_points = [v for v in np.arange(0, 14500, 1000)[1:] if min_dist(fixed_points, v) > 200] - vlist = sorted(list([*fixed_points, *regular_points])) + vlist = sorted([*fixed_points, *regular_points]) v_inner = 0.0 # velocity at inner boundary of cell m_tot = 0.0 @@ -92,7 +91,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: dfmodel.loc[cellid] = [cellid + 1, v_outer, math.log10(rho), *radioabundances] dfelabundances.loc[cellid] = [cellid + 1, *abundances[1:31]] - r_inner, r_outer = [(v * u.km / u.s * t200 * 200 * u.day).to("cm").value for v in [v_inner, v_outer]] + r_inner, r_outer = ((v * u.km / u.s * t200 * 200 * u.day).to("cm").value for v in [v_inner, v_outer]) vol_shell = 4 * math.pi / 3 * (r_outer**3 - r_inner**3) m_shell = rho * vol_shell / u.solMass.to("g") @@ -102,7 +101,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: print(f"M_tot = {m_tot:.3f} solMass") at.inputmodel.save_modeldata(dfmodel, t_model_init_days, os.path.join(args.outputpath, "model.txt")) - at.inputmodel.save_initialabundances(dfelabundances, os.path.join(args.outputpath, "abundances.txt")) + at.inputmodel.save_initelemabundances(dfelabundances, os.path.join(args.outputpath, "abundances.txt")) if __name__ == "__main__": diff --git a/artistools/inputmodel/describeinputmodel.py b/artistools/inputmodel/describeinputmodel.py old mode 100644 new mode 100755 index 32a107111..d0cfa92cf --- a/artistools/inputmodel/describeinputmodel.py +++ b/artistools/inputmodel/describeinputmodel.py @@ -9,8 +9,6 @@ import artistools as at -# import pandas as pd - def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-inputfile", "-i", default=Path(), help="Path of input file or folder containing model.txt") @@ -45,9 +43,10 @@ def main(args=None, argsraw=None, **kwargs): if args.noisotopes: args.getelemabundances = True - dfmodel, t_model_init_days, vmax = at.inputmodel.get_modeldata_tuple( - args.inputfile, get_elemabundances=args.getelemabundances, printwarningsonly=False + dfmodel, modelmeta = at.inputmodel.get_modeldata( + args.inputfile, get_elemabundances=args.getelemabundances, printwarningsonly=False, dtype_backend="pyarrow" ) + t_model_init_days, vmax = modelmeta["t_model_init_days"], modelmeta["vmax_cmps"] t_model_init_seconds = t_model_init_days * 24 * 60 * 60 print(f"Model is defined at {t_model_init_days} days ({t_model_init_seconds:.4f} seconds)") @@ -70,7 +69,7 @@ def main(args=None, argsraw=None, **kwargs): mgi = int(args.cell) if mgi >= 0: print(f"Selected single cell mgi {mgi}:") - dfmodel.query("inputcellid == (@mgi + 1)", inplace=True) + dfmodel = dfmodel.query("inputcellid == (@mgi + 1)") print(dfmodel.iloc[0]) mass_msun_rho = dfmodel["cellmass_grams"].sum() / 1.989e33 @@ -132,7 +131,7 @@ def sortkey(tup_species_mass_g: tuple[str, float]): elem_mass = speciesmasses.get(elsymb, 0.0) if elem_mass > 0.0: strcomment += f" ({mass_g / elem_mass * 100:6.2f}% of {elsymb} element mass)" - if mass_g > elem_mass * (1.0 + 1e-10): + if mass_g > elem_mass * (1.0 + 1e-5): strcomment += " ERROR! isotope sum is greater than element abundance" zstr = f"Z={atomic_number}" print(f"{zstr:>5} {species:9s} {species_mass_msun:.3e} Msun massfrac {massfrac:.3e}{strcomment}") diff --git a/artistools/inputmodel/downscale3dgrid.py b/artistools/inputmodel/downscale3dgrid.py index 1d08818eb..885ca0183 100644 --- a/artistools/inputmodel/downscale3dgrid.py +++ b/artistools/inputmodel/downscale3dgrid.py @@ -32,7 +32,7 @@ def make_downscaled_3d_grid(modelpath, inputgridsize=200, outputgridsize=50, plo abread = np.zeros(31) print("reading abundance file") - with open(abundancefile, "r") as sourceabundancefile: + with open(abundancefile) as sourceabundancefile: for z in range(0, grid): for y in range(0, grid): for x in range(0, grid): @@ -40,7 +40,7 @@ def make_downscaled_3d_grid(modelpath, inputgridsize=200, outputgridsize=50, plo abund[x, y, z] = abread print("reading model file") - with open(modelfile, "r") as sourcemodelfile: + with open(modelfile) as sourcemodelfile: x = sourcemodelfile.readline() t_model = sourcemodelfile.readline() vmax = sourcemodelfile.readline() diff --git a/artistools/inputmodel/energyinputfiles.py b/artistools/inputmodel/energyinputfiles.py index f40012bee..60647a2b6 100644 --- a/artistools/inputmodel/energyinputfiles.py +++ b/artistools/inputmodel/energyinputfiles.py @@ -8,8 +8,6 @@ from scipy import integrate import artistools as at -import artistools.inputmodel - DAY = 86400 # day in seconds MSUN = 1.989e33 # solar mass in grams @@ -84,9 +82,7 @@ def define_heating_rate(): E_tot = integrate.trapezoid(y=qdot, x=times) # ergs/s/g # print("Etot per gram", E_tot, E_tot*1.989e33*0.01) - import scipy.integrate - - cumulative_integrated_energy = scipy.integrate.cumulative_trapezoid(y=qdot, x=times) + cumulative_integrated_energy = integrate.cumulative_trapezoid(y=qdot, x=times) cumulative_integrated_energy = np.insert(cumulative_integrated_energy, 0, 0) rate = cumulative_integrated_energy / E_tot @@ -123,7 +119,7 @@ def define_heating_rate(): def energy_from_rprocess_calculation(energy_thermo_data, get_rate=True): index_time_greaterthan = energy_thermo_data[energy_thermo_data["time/s"] > 1e7].index # 1e7 seconds = 116 days - energy_thermo_data.drop(index_time_greaterthan, inplace=True) + energy_thermo_data = energy_thermo_data.drop(index_time_greaterthan) # print("Dropping times later than 116 days") skipfirstnrows = 0 # not sure first values look sensible -- check this @@ -135,9 +131,7 @@ def energy_from_rprocess_calculation(energy_thermo_data, get_rate=True): if get_rate: print(f"E_tot {E_tot} erg/g") - import scipy.integrate - - cumulative_integrated_energy = scipy.integrate.cumulative_trapezoid(y=qdot, x=times) + cumulative_integrated_energy = integrate.cumulative_trapezoid(y=qdot, x=times) cumulative_integrated_energy = np.insert(cumulative_integrated_energy, 0, 0) rate = cumulative_integrated_energy / E_tot @@ -147,8 +141,7 @@ def energy_from_rprocess_calculation(energy_thermo_data, get_rate=True): return times_and_rate, E_tot - else: - return E_tot + return E_tot def get_rprocess_calculation_files(path_to_rprocess_calculation, interpolate_trajectories=False, thermalisation=False): @@ -186,7 +179,7 @@ def get_rprocess_calculation_files(path_to_rprocess_calculation, interpolate_tra interpolated_trajectories["mean"] = interpolated_trajectories.iloc[:, 1:].mean(axis=1) index_time_lessthan = interpolated_trajectories[interpolated_trajectories["time/s"] < 1.1e-1].index - interpolated_trajectories.drop(index_time_lessthan, inplace=True) + interpolated_trajectories = interpolated_trajectories.drop(index_time_lessthan) interpolated_trajectories.to_csv(path_to_rprocess_calculation / "interpolatedQdot.dat", sep=" ", index=False) print(f"sum etot {sum(trajectory_E_tot)}") @@ -248,3 +241,22 @@ def plot_energy_rate(modelpath): plt.plot( times_and_rate["times"], np.array(times_and_rate["nuclear_heating_power"]) * Mtot_grams, color="k", zorder=10 ) + + +def get_etot_fromfile(modelpath): + energydistribution_data = pd.read_csv( + Path(modelpath) / "energydistribution.txt", + skiprows=1, + delim_whitespace=True, + header=None, + names=["cellid", "cell_energy"], + ) + etot = energydistribution_data["cell_energy"].sum() + return etot, energydistribution_data + + +def get_energy_rate_fromfile(modelpath): + energyrate_data = pd.read_csv( + Path(modelpath) / "energyrate.txt", skiprows=1, delim_whitespace=True, header=None, names=["times", "rate"] + ) + return energyrate_data diff --git a/artistools/inputmodel/fromcmfgen/__init__.py b/artistools/inputmodel/fromcmfgen/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/artistools/inputmodel/fromcmfgen/convert_to_artis_neartimezero.py b/artistools/inputmodel/fromcmfgen/convert_to_artis_neartimezero.py index f7b29758a..56148b2e6 100755 --- a/artistools/inputmodel/fromcmfgen/convert_to_artis_neartimezero.py +++ b/artistools/inputmodel/fromcmfgen/convert_to_artis_neartimezero.py @@ -6,9 +6,8 @@ from math import exp import numpy as np -from rd_cmfgen import rd_sn_hydro_data -# import math +from .rd_cmfgen import rd_sn_hydro_data msun = 1.989e33 @@ -224,7 +223,7 @@ def main(): co56frac = a["isofrac"][i, 89] cr48frac = a["isofrac"][i, 75] v48frac = a["isofrac"][i, 49] - strout = "{0:4d} {1:1.7e} {2:1.7e} {3:1.7e} {4:1.7e} {5:1.7e} {6:1.7e} {7:1.7e}\n".format( + strout = "{:4d} {:1.7e} {:1.7e} {:1.7e} {:1.7e} {:1.7e} {:1.7e} {:1.7e}\n".format( i + 1, vel, rho, igefrac, ni56frac, co56frac, cr48frac, v48frac ) f.write(strout) diff --git a/artistools/inputmodel/fromcmfgen/rd_cmfgen.py b/artistools/inputmodel/fromcmfgen/rd_cmfgen.py index 50cdb4e4a..9068e4c63 100644 --- a/artistools/inputmodel/fromcmfgen/rd_cmfgen.py +++ b/artistools/inputmodel/fromcmfgen/rd_cmfgen.py @@ -1,17 +1,13 @@ +# ruff: noqa """various functions to read/write CMFGEN input/output files: rd_nuc_decay_data rd_sn_hydro_data """ -# import os import sys import numpy as np -# import re - -# from pdb import set_trace as stop - # constants DAY2SEC = 86400.0 MEV2ERG = 1.60217733e-6 @@ -23,7 +19,7 @@ def rd_nuc_decay_data(file, quiet=False): """ # open file - with open(file, "r") as f: + with open(file) as f: # read in header while True: line = f.readline() @@ -60,7 +56,7 @@ def rd_nuc_decay_data(file, quiet=False): isospec.append(linearr[0]) amu[i] = float(linearr[1]) aiso[i] = np.rint(amu[i]) - stable.append(True) if linearr[2] == "s" else stable.append(False) + stable.append(linearr[2] == "s") if not quiet: print("INFO - Read in isotope information") @@ -139,7 +135,7 @@ def rd_sn_hydro_data(file, ncol=8, reverse=False, quiet=False): MAX_POP_DIFF = 1e-5 # maximum absolute difference between sum(isofrac) and corresponding specfrac # open file - with open(file, "r") as f: + with open(file) as f: # read in header okhdr = 0 nd, nspec, niso = 0, 0, 0 @@ -179,7 +175,7 @@ def rd_sn_hydro_data(file, ncol=8, reverse=False, quiet=False): kappa = np.zeros(nd) # mass absorption coefficient (cm^2/g) okhydro = 0 while okhydro == 0: - while line == "": + while not line: line = f.readline() if "Radius grid" in line: rad = np.fromfile(f, count=nd, sep=" ", dtype=float) @@ -288,7 +284,7 @@ def rd_sn_hydro_data(file, ncol=8, reverse=False, quiet=False): absdiff = np.abs(specfrac[:, spec.index(s)] - sumisofrac) relabsdiff = absdiff / sumisofrac if np.max(relabsdiff) > MAX_POP_DIFF: - sys.exit("ERROR - Maximum absolute difference > MAX_POP_DIFF for species {0:s}".format(s)) + sys.exit("ERROR - Maximum absolute difference > MAX_POP_DIFF for species {:s}".format(s)) # reversed vectors if reverse=True if reverse: diff --git a/artistools/inputmodel/fullymixed.py b/artistools/inputmodel/fullymixed.py index 9769564ed..1834a6483 100755 --- a/artistools/inputmodel/fullymixed.py +++ b/artistools/inputmodel/fullymixed.py @@ -6,11 +6,6 @@ import artistools as at -# import math -# import os.path -# import numpy as np -# import pandas as pd - def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-scalefactor", "-s", default=0.5, help="Kinetic energy scale factor") @@ -19,12 +14,11 @@ def addargs(parser: argparse.ArgumentParser) -> None: def eval_mshell(dfmodel, t_model_init_seconds): - dfmodel.eval( + dfmodel = dfmodel.eval( ( "cellmass_grams = 10 ** logrho * 4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" "* (1e5 * @t_model_init_seconds) ** 3" ), - inplace=True, ) @@ -41,7 +35,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: dfmodel, t_model_init_days, _ = at.inputmodel.get_modeldata_tuple(args.inputpath) print("Read model.txt") - dfelabundances = at.inputmodel.get_initialabundances(args.inputpath) + dfelabundances = at.inputmodel.get_initelemabundances(args.inputpath) print("Read abundances.txt") t_model_init_seconds = t_model_init_days * 24 * 60 * 60 @@ -57,13 +51,13 @@ def main(args=None, argsraw=None, **kwargs) -> None: integrated_mass_grams = (dfmodel[column_name] * dfmodel.cellmass_grams).sum() global_massfrac = integrated_mass_grams / model_mass_grams print(f"{column_name:>13s}: {global_massfrac:.3f} ({integrated_mass_grams * u.g.to('solMass'):.3f} Msun)") - dfmodel.eval(f"{column_name} = {global_massfrac}", inplace=True) + dfmodel = dfmodel.eval(f"{column_name} = {global_massfrac}") for column_name in [x for x in dfelabundances.columns if x.startswith("X_")]: integrated_mass_grams = (dfelabundances[column_name] * dfmodel.cellmass_grams).sum() global_massfrac = integrated_mass_grams / model_mass_grams print(f"{column_name:>13s}: {global_massfrac:.3f} ({integrated_mass_grams * u.g.to('solMass'):.3f} Msun)") - dfelabundances.eval(f"{column_name} = {global_massfrac}", inplace=True) + dfelabundances = dfelabundances.eval(f"{column_name} = {global_massfrac}") print(dfmodel) print(dfelabundances) @@ -73,7 +67,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: print(f"Saved {modeloutfilename}") abundoutfilename = "abundances_fullymixed.txt" - at.inputmodel.save_initialabundances(dfelabundances, Path(args.outputpath, abundoutfilename)) + at.inputmodel.save_initelemabundances(dfelabundances, Path(args.outputpath, abundoutfilename)) print(f"Saved {abundoutfilename}") diff --git a/artistools/inputmodel/inputmodel_misc.py b/artistools/inputmodel/inputmodel_misc.py index 070f94ad6..ae93d35e8 100644 --- a/artistools/inputmodel/inputmodel_misc.py +++ b/artistools/inputmodel/inputmodel_misc.py @@ -2,44 +2,47 @@ import gc import math import os.path +import pickle import time +from collections import defaultdict from collections.abc import Sequence from functools import lru_cache from pathlib import Path from typing import Any +from typing import Callable +from typing import Literal from typing import Optional from typing import Union import numpy as np import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq import artistools as at -# from collections import namedtuple - -@lru_cache(maxsize=8) -def read_modelfile( +def read_modelfile_text( filename: Union[Path, str], - dimensions: Optional[int] = None, printwarningsonly: bool = False, getheadersonly: bool = False, - skipabundancecolumns: bool = False, + skipnuclidemassfraccolumns: bool = False, + dtype_backend: Literal["pyarrow", "numpy_nullable"] = "numpy_nullable", ) -> tuple[pd.DataFrame, dict[str, Any]]: """ Read an artis model.txt file containing cell velocities, density, and abundances of radioactive nuclides. """ - assert dimensions in [1, 3, None] onelinepercellformat = None modelmeta: dict[str, Any] = {"headercommentlines": []} modelpath = Path(filename).parent - print(f"Reading {filename}") + if not printwarningsonly: + print(f"Reading {filename}") numheaderrows = 0 - with at.misc.zopen(filename, "rt") as fmodel: + with at.zopen(filename) as fmodel: line = "#" while line.startswith("#"): line = fmodel.readline() @@ -47,33 +50,38 @@ def read_modelfile( modelmeta["headercommentlines"].append(line.removeprefix("#").removeprefix(" ").removesuffix("\n")) numheaderrows += 1 - modelcellcount = int(line) - t_model_init_days = float(fmodel.readline()) + if len(line.strip().split(" ")) == 2: + print(" detected 2D model file") + modelmeta["dimensions"] = 2 + ncoordgrid_r, ncoordgrid_z = (int(n) for n in line.strip().split(" ")) + modelmeta["ncoordgrid_r"] = ncoordgrid_r + modelmeta["ncoordgrid_z"] = ncoordgrid_z + modelcellcount = ncoordgrid_r * ncoordgrid_z + else: + modelcellcount = int(line) + + modelmeta["npts_model"] = modelcellcount + modelmeta["t_model_init_days"] = float(fmodel.readline()) numheaderrows += 2 - t_model_init_seconds = t_model_init_days * 24 * 60 * 60 + t_model_init_seconds = modelmeta["t_model_init_days"] * 24 * 60 * 60 filepos = fmodel.tell() - # if the next line is a single float then the model is 3D + # if the next line is a single float then the model is 2D or 3D (vmax) try: - vmax_cmps = float(fmodel.readline()) # velocity max in cm/s - xmax_tmodel = vmax_cmps * t_model_init_seconds # xmax = ymax = zmax + modelmeta["vmax_cmps"] = float(fmodel.readline()) # velocity max in cm/s + xmax_tmodel = modelmeta["vmax_cmps"] * t_model_init_seconds # xmax = ymax = zmax numheaderrows += 1 - if dimensions is None: + if "dimensions" not in modelmeta: if not printwarningsonly: print(" detected 3D model file") - dimensions = 3 - elif dimensions != 3: - print(f" {dimensions} were specified but file appears to be 3D") - assert False + modelmeta["dimensions"] = 3 except ValueError: - if dimensions is None: + assert modelmeta.get("dimensions", -1) != 2 # 2D model should have vmax line here + if "dimensions" not in modelmeta: if not printwarningsonly: print(" detected 1D model file") - dimensions = 1 - elif dimensions != 1: - print(f" {dimensions} were specified but file appears to be 1D") - assert False + modelmeta["dimensions"] = 1 fmodel.seek(filepos) # undo the readline() and go back @@ -90,7 +98,7 @@ def read_modelfile( ncols_line_even = len(data_line_even) if columns is None: - if dimensions == 1: + if modelmeta["dimensions"] == 1: columns = [ "inputcellid", "velocity_outer", @@ -104,7 +112,22 @@ def read_modelfile( "X_Co57", ][:ncols_line_even] - elif dimensions == 3: + elif modelmeta["dimensions"] == 2: + columns = [ + "inputcellid", + "pos_r_mid", + "pos_z_mid", + "rho", + "X_Fegroup", + "X_Ni56", + "X_Co56", + "X_Fe52", + "X_Cr48", + "X_Ni57", + "X_Co57", + ][:ncols_line_even] + + elif modelmeta["dimensions"] == 3: columns = [ "inputcellid", "inputpos_a", @@ -119,55 +142,73 @@ def read_modelfile( "X_Ni57", "X_Co57", ][:ncols_line_even] - else: - # TODO: 2D case - assert False + + assert columns is not None if ncols_line_even == len(columns): - print(" model file is one line per cell") + if not printwarningsonly: + print(" model file is one line per cell") ncols_line_odd = 0 onelinepercellformat = True else: - print(" model file format is two lines per cell") + if not printwarningsonly: + print(" model file format is two lines per cell") # columns split over two lines ncols_line_odd = len(fmodel.readline().split()) assert (ncols_line_even + ncols_line_odd) == len(columns) onelinepercellformat = False - if skipabundancecolumns: - print(" skipping abundance columns in model.txt") - if dimensions == 1: + if skipnuclidemassfraccolumns: + if not printwarningsonly: + print(" skipping nuclide abundance columns in model") + if modelmeta["dimensions"] == 1: ncols_line_even = 3 - elif dimensions == 3: + elif modelmeta["dimensions"] == 2: + ncols_line_even = 4 + elif modelmeta["dimensions"] == 3: ncols_line_even = 5 ncols_line_odd = 0 - if dimensions == 3: + if modelmeta["dimensions"] == 3: # number of grid cell steps along an axis (same for xyz) ncoordgridx = int(round(modelcellcount ** (1.0 / 3.0))) ncoordgridy = int(round(modelcellcount ** (1.0 / 3.0))) ncoordgridz = int(round(modelcellcount ** (1.0 / 3.0))) + modelmeta["ncoordgridx"] = ncoordgridx + modelmeta["ncoordgridy"] = ncoordgridy + modelmeta["ncoordgridz"] = ncoordgridz + if ncoordgridx == ncoordgridy == ncoordgridz: + modelmeta["ncoordgrid"] = ncoordgridx assert (ncoordgridx * ncoordgridy * ncoordgridz) == modelcellcount nrows_read = 1 if getheadersonly else modelcellcount skiprows: Union[list, int, None] - if onelinepercellformat: - skiprows = numheaderrows + + skiprows = ( + numheaderrows + if onelinepercellformat + else [ + x + for x in range(numheaderrows + modelcellcount * 2) + if x < numheaderrows or (x - numheaderrows - 1) % 2 == 0 + ] + ) + + dtypes: defaultdict[str, Union[Callable, str]] + if dtype_backend == "pyarrow": + dtypes = defaultdict(lambda: "float32[pyarrow]") + dtypes["inputcellid"] = "int32[pyarrow]" + dtypes["tracercount"] = "int32[pyarrow]" else: - # skip the odd rows for the first read in - skiprows = list( - [ - x - for x in range(numheaderrows + modelcellcount * 2) - if x < numheaderrows or (x - numheaderrows - 1) % 2 == 0 - ] - ) + dtypes = defaultdict(lambda: "float32") + dtypes["inputcellid"] = "int32" + dtypes["tracercount"] = "int32" # each cell takes up two lines in the model file - dfmodel = pd.read_table( - filename, + dfmodel = pd.read_csv( + at.zopen(filename), sep=r"\s+", engine="c", header=None, @@ -175,25 +216,27 @@ def read_modelfile( names=columns[:ncols_line_even], usecols=columns[:ncols_line_even], nrows=nrows_read, + dtype=dtypes, + dtype_backend=dtype_backend, ) if ncols_line_odd > 0 and not onelinepercellformat: # read in the odd rows and merge dataframes - skipevenrows = list( - [ - x - for x in range(numheaderrows + modelcellcount * 2) - if x < numheaderrows or (x - numheaderrows - 1) % 2 == 1 - ] - ) - dfmodeloddlines = pd.read_table( - filename, + skipevenrows = [ + x + for x in range(numheaderrows + modelcellcount * 2) + if x < numheaderrows or (x - numheaderrows - 1) % 2 == 1 + ] + dfmodeloddlines = pd.read_csv( + at.zopen(filename), sep=r"\s+", engine="c", header=None, skiprows=skipevenrows, names=columns[ncols_line_even:], nrows=nrows_read, + dtype=dtypes, + dtype_backend=dtype_backend, ) assert len(dfmodel) == len(dfmodeloddlines) dfmodel = dfmodel.merge(dfmodeloddlines, left_index=True, right_index=True) @@ -207,24 +250,24 @@ def read_modelfile( dfmodel.index.name = "cellid" # dfmodel.drop('inputcellid', axis=1, inplace=True) - if dimensions == 1: - dfmodel["velocity_inner"] = np.concatenate([[0.0], dfmodel["velocity_outer"].values[:-1]]) - dfmodel.eval( - ( - "cellmass_grams = 10 ** logrho * 4. / 3. * 3.14159265 * (velocity_outer ** 3 - velocity_inner ** 3)" - "* (1e5 * @t_model_init_seconds) ** 3" - ), - inplace=True, + if modelmeta["dimensions"] == 1: + dfmodel["velocity_inner"] = np.concatenate([[0.0], dfmodel["velocity_outer"].to_numpy()[:-1]]) + dfmodel["cellmass_grams"] = ( + 10 ** dfmodel["logrho"] + * (4.0 / 3.0) + * 3.14159265 + * (dfmodel["velocity_outer"] ** 3 - dfmodel["velocity_inner"] ** 3) + * (1e5 * t_model_init_seconds) ** 3 ) - vmax_cmps = dfmodel.velocity_outer.max() * 1e5 + modelmeta["vmax_cmps"] = dfmodel.velocity_outer.max() * 1e5 - elif dimensions == 3: - wid_init = at.misc.get_wid_init_at_tmodel(modelpath, modelcellcount, t_model_init_days, xmax_tmodel) + elif modelmeta["dimensions"] == 3: + wid_init = at.get_wid_init_at_tmodel(modelpath, modelcellcount, modelmeta["t_model_init_days"], xmax_tmodel) modelmeta["wid_init"] = wid_init - dfmodel.eval("cellmass_grams = rho * @wid_init ** 3", inplace=True) + dfmodel["cellmass_grams"] = dfmodel["rho"] * wid_init**3 - dfmodel.rename(columns={"pos_x": "pos_x_min", "pos_y": "pos_y_min", "pos_z": "pos_z_min"}, inplace=True) - if "pos_x_min" in dfmodel.columns: + dfmodel = dfmodel.rename(columns={"pos_x": "pos_x_min", "pos_y": "pos_y_min", "pos_z": "pos_z_min"}) + if "pos_x_min" in dfmodel.columns and not printwarningsonly: print(" model cell positions are defined in the header") elif not getheadersonly: @@ -267,20 +310,15 @@ def vectormatch(vec1, vec2): assert posmatch_xyz != posmatch_zyx # one option must match if posmatch_xyz: print(" model cell positions are consistent with x-y-z column order") - dfmodel.rename( + dfmodel = dfmodel.rename( columns={"inputpos_a": "pos_x_min", "inputpos_b": "pos_y_min", "inputpos_c": "pos_z_min"}, - inplace=True, ) if posmatch_zyx: print(" cell positions are consistent with z-y-x column order") - dfmodel.rename( + dfmodel = dfmodel.rename( columns={"inputpos_a": "pos_z_min", "inputpos_b": "pos_y_min", "inputpos_c": "pos_x_min"}, - inplace=True, ) - modelmeta["t_model_init_days"] = t_model_init_days - modelmeta["dimensions"] = dimensions - modelmeta["vmax_cmps"] = vmax_cmps modelmeta["modelcellcount"] = modelcellcount return dfmodel, modelmeta @@ -288,32 +326,31 @@ def vectormatch(vec1, vec2): def get_modeldata( inputpath: Union[Path, str] = Path(), - dimensions: Optional[int] = None, get_elemabundances: bool = False, derived_cols: Optional[Sequence[str]] = None, printwarningsonly: bool = False, getheadersonly: bool = False, - skipabundancecolumns: bool = False, + skipnuclidemassfraccolumns: bool = False, + dtype_backend: Literal["pyarrow", "numpy_nullable"] = "numpy_nullable", ) -> tuple[pd.DataFrame, dict[str, Any]]: """ - Read an artis model.txt file containing cell velocities, density, and abundances of radioactive nuclides. + Read an artis model.txt file containing cell velocities, densities, and mass fraction abundances of radioactive nuclides. Parameters: - inputpath: either a path to model.txt file, or a folder containing model.txt - - dimensions: number of dimensions in input file, or None for automatic - get_elemabundances: also read elemental abundances (abundances.txt) and merge with the output DataFrame return dfmodel, modelmeta - dfmodel: a pandas DataFrame with a row for each model grid cell - - modelmeta is a dictionary of model properties, such as t_model_init_days, vmax_cmps, dimensions, etc + - modelmeta: a dictionary of input model parameters, with keys such as t_model_init_days, vmax_cmps, dimensions, etc. """ inputpath = Path(inputpath) if inputpath.is_dir(): modelpath = inputpath - filename = at.firstexisting("model.txt", path=inputpath, tryzipped=True) + filename = at.firstexisting("model.txt", folder=inputpath, tryzipped=True) elif inputpath.is_file(): # passed in a filename instead of the modelpath filename = inputpath modelpath = Path(inputpath).parent @@ -324,33 +361,80 @@ def get_modeldata( else: raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), inputpath) - # this inner call gets LRU cached, so that we can reuse the cache when changing derived columns, or adding elemental abundances - dfmodel, modelmeta = read_modelfile( - filename=filename, - dimensions=dimensions, - printwarningsonly=printwarningsonly, - getheadersonly=getheadersonly, - skipabundancecolumns=skipabundancecolumns, - ) + dfmodel = None + filenameparquet = at.stripallsuffixes(Path(filename)).with_suffix(".txt.parquet") - dfmodel = dfmodel.copy() + source_textfile_details = {"st_size": filename.stat().st_size, "st_mtime": filename.stat().st_mtime} + + if filenameparquet.is_file() and not getheadersonly: + if not printwarningsonly: + print(f" reading data table from {filenameparquet}") + + pqmetadata = pq.read_metadata(filenameparquet) + if ( + b"artismodelmeta" not in pqmetadata.metadata + or b"source_textfile_details" not in pqmetadata.metadata + or pickle.dumps(source_textfile_details) != pqmetadata.metadata[b"source_textfile_details"] + ): + print(f" text source {filename} doesn't match file header of {filenameparquet}. Removing parquet file") + filenameparquet.unlink(missing_ok=True) + else: + modelmeta = pickle.loads(pqmetadata.metadata[b"artismodelmeta"]) + + columns = ( + [col for col in pqmetadata.schema.names if not col.startswith("X_")] + if skipnuclidemassfraccolumns + else None + ) + dfmodel = pd.read_parquet( + filenameparquet, + columns=columns, + dtype_backend=dtype_backend, + ) + + if dfmodel is None: + skipnuclidemassfraccolumns = False + dfmodel, modelmeta = read_modelfile_text( + filename=filename, + printwarningsonly=printwarningsonly, + getheadersonly=getheadersonly, + skipnuclidemassfraccolumns=skipnuclidemassfraccolumns, + dtype_backend=dtype_backend, + ) + + if len(dfmodel) > 1000 and not getheadersonly and not skipnuclidemassfraccolumns: + print(f"Saving {filenameparquet}") + patable = pa.Table.from_pandas(dfmodel) + + custom_metadata = { + b"source_textfile_details": pickle.dumps(source_textfile_details), + b"artismodelmeta": pickle.dumps(modelmeta), + } + merged_metadata = {**custom_metadata, **(patable.schema.metadata or {})} + patable = patable.replace_schema_metadata(merged_metadata) + pq.write_table(patable, filenameparquet, compression="ZSTD") + # dfmodel.to_parquet(filenameparquet, compression="zstd") + print(" Done.") if get_elemabundances: - if dimensions == 3: - print("Getting abundances") - abundancedata = get_initialabundances(modelpath) + abundancedata = get_initelemabundances( + modelpath, dtype_backend=dtype_backend, printwarningsonly=printwarningsonly + ) dfmodel = dfmodel.merge(abundancedata, how="inner", on="inputcellid") if derived_cols: - add_derived_cols_to_modeldata( + dfmodel = add_derived_cols_to_modeldata( dfmodel=dfmodel, derived_cols=derived_cols, - dimensions=dimensions, + dimensions=modelmeta["dimensions"], t_model_init_seconds=modelmeta["t_model_init_days"] * 86400.0, wid_init=modelmeta.get("wid_init", None), modelpath=modelpath, ) + if len(dfmodel) > 100000: + dfmodel.info(verbose=False, memory_usage="deep") + return dfmodel, modelmeta @@ -379,40 +463,38 @@ def add_derived_cols_to_modeldata( if dimensions == 3: if "velocity" in derived_cols or "vel_min" in derived_cols: assert t_model_init_seconds is not None - dfmodel["vel_x_min"] = dfmodel["pos_x_min"] / t_model_init_seconds - dfmodel["vel_y_min"] = dfmodel["pos_y_min"] / t_model_init_seconds - dfmodel["vel_z_min"] = dfmodel["pos_z_min"] / t_model_init_seconds + for ax in ["x", "y", "z"]: + dfmodel[f"vel_{ax}_min"] = dfmodel[f"pos_{ax}_min"] / t_model_init_seconds if "velocity" in derived_cols or "vel_max" in derived_cols: assert t_model_init_seconds is not None - dfmodel["vel_x_max"] = (dfmodel["pos_x_min"] + wid_init) / t_model_init_seconds - dfmodel["vel_y_max"] = (dfmodel["pos_y_min"] + wid_init) / t_model_init_seconds - dfmodel["vel_z_max"] = (dfmodel["pos_z_min"] + wid_init) / t_model_init_seconds + for ax in ["x", "y", "z"]: + dfmodel[f"vel_{ax}_max"] = (dfmodel[f"pos_{ax}_min"] + wid_init) / t_model_init_seconds - if any([col in derived_cols for col in ["velocity", "vel_mid", "vel_mid_radial"]]): + if any(col in derived_cols for col in ["velocity", "vel_mid", "vel_r_mid"]): assert wid_init is not None assert t_model_init_seconds is not None - dfmodel["vel_x_mid"] = (dfmodel["pos_x_min"] + (0.5 * wid_init)) / t_model_init_seconds - dfmodel["vel_y_mid"] = (dfmodel["pos_y_min"] + (0.5 * wid_init)) / t_model_init_seconds - dfmodel["vel_z_mid"] = (dfmodel["pos_z_min"] + (0.5 * wid_init)) / t_model_init_seconds + for ax in ["x", "y", "z"]: + dfmodel[f"vel_{ax}_mid"] = (dfmodel[f"pos_{ax}_min"] + (0.5 * wid_init)) / t_model_init_seconds - dfmodel.eval("vel_mid_radial = sqrt(vel_x_mid ** 2 + vel_y_mid ** 2 + vel_z_mid ** 2)", inplace=True) + dfmodel["vel_r_mid"] = np.sqrt( + dfmodel["vel_x_mid"] ** 2 + dfmodel["vel_y_mid"] ** 2 + dfmodel["vel_z_mid"] ** 2 + ) if dimensions == 3 and "pos_mid" in derived_cols or "angle_bin" in derived_cols: assert wid_init is not None - dfmodel["pos_x_mid"] = dfmodel["pos_x_min"] + (0.5 * wid_init) - dfmodel["pos_y_mid"] = dfmodel["pos_y_min"] + (0.5 * wid_init) - dfmodel["pos_z_mid"] = dfmodel["pos_z_min"] + (0.5 * wid_init) + for ax in ["x", "y", "z"]: + dfmodel[f"pos_{ax}_mid"] = dfmodel[f"pos_{ax}_min"] + (0.5 * wid_init) if "logrho" in derived_cols and "logrho" not in dfmodel.columns: - dfmodel.eval("logrho = log10(rho)", inplace=True) + dfmodel["logrho"] = np.log10(dfmodel["rho"]) if "rho" in derived_cols and "rho" not in dfmodel.columns: - dfmodel.eval("rho = 10**logrho", inplace=True) + dfmodel["rho"] = 10 ** dfmodel["logrho"] if "angle_bin" in derived_cols: assert modelpath is not None - get_cell_angle(dfmodel, modelpath) + dfmodel = get_cell_angle(dfmodel, modelpath) # if "Ye" in derived_cols and os.path.isfile(modelpath / "Ye.txt"): # dfmodel["Ye"] = at.inputmodel.opacityinputfile.get_Ye_from_file(modelpath) @@ -450,7 +532,6 @@ def get_mean_cell_properties_of_angle_bin( get_cell_angle(dfmodeldata, modelpath) dfmodeldata["rho"][dfmodeldata["rho"] == 0] = None - dfmodeldata["rho"] cell_velocities = np.unique(dfmodeldata["vel_x_min"].values) cell_velocities = cell_velocities[cell_velocities >= 0] @@ -476,17 +557,17 @@ def get_mean_cell_properties_of_angle_bin( # get cells with bin number dfanglebin = dfmodeldata.query("cos_bin == @cos_bin_number", inplace=False) - binned = pd.cut(dfanglebin["vel_mid_radial"], velocity_bins, labels=False, include_lowest=True) + binned = pd.cut(dfanglebin["vel_r_mid"], velocity_bins, labels=False, include_lowest=True) i = 0 for binindex, mean_rho in dfanglebin.groupby(binned)["rho"].mean().iteritems(): i += 1 mean_bin_properties[bin_number]["mean_rho"][binindex] += mean_rho i = 0 - if "Ye" in dfmodeldata.keys(): + if "Ye" in dfmodeldata: for binindex, mean_Ye in dfanglebin.groupby(binned)["Ye"].mean().iteritems(): i += 1 mean_bin_properties[bin_number]["mean_Ye"][binindex] += mean_Ye - if "Q" in dfmodeldata.keys(): + if "Q" in dfmodeldata: for binindex, mean_Q in dfanglebin.groupby(binned)["Q"].mean().iteritems(): i += 1 mean_bin_properties[bin_number]["mean_Q"][binindex] += mean_Q @@ -524,9 +605,9 @@ def get_3d_model_data_merged_model_and_abundances_minimal(args): """Get 3D data without generating all the extra columns in standard routine. Needed for large (eg. 200^3) models""" model = get_3d_modeldata_minimal(args.modelpath) - abundances = get_initialabundances(args.modelpath[0]) + abundances = get_initelemabundances(args.modelpath[0]) - with open(os.path.join(args.modelpath[0], "model.txt"), "r") as fmodelin: + with open(os.path.join(args.modelpath[0], "model.txt")) as fmodelin: fmodelin.readline() # npts_model3d args.t_model = float(fmodelin.readline()) # days args.vmax = float(fmodelin.readline()) # v_max in [cm/s] @@ -562,7 +643,7 @@ def get_3d_modeldata_minimal(modelpath) -> pd.DataFrame: "X_Fe52", "X_Cr48", ] - model = pd.DataFrame(model.values.reshape(-1, 10)) + model = pd.DataFrame(model.to_numpy().reshape(-1, 10)) model.columns = columns print("model.txt memory usage:") @@ -579,6 +660,8 @@ def save_modeldata( dimensions: Optional[int] = None, headercommentlines: Optional[list[str]] = None, modelmeta: Optional[dict[str, Any]] = None, + twolinespercell: bool = False, + float_format: str = ".4e", ) -> None: """Save a pandas DataFrame and snapshot time into ARTIS model.txt""" if modelmeta: @@ -609,7 +692,7 @@ def save_modeldata( if dimensions == 1: standardcols = ["inputcellid", "velocity_outer", "logrho", "X_Fegroup", "X_Ni56", "X_Co56", "X_Fe52", "X_Cr48"] elif dimensions == 3: - dfmodel.rename(columns={"gridindex": "inputcellid"}, inplace=True) + dfmodel = dfmodel.rename(columns={"gridindex": "inputcellid"}) griddimension = int(round(len(dfmodel) ** (1.0 / 3.0))) print(f" grid size: {len(dfmodel)} ({griddimension}^3)") assert griddimension**3 == len(dfmodel) @@ -634,24 +717,21 @@ def save_modeldata( if "X_Co57" in dfmodel.columns: standardcols.append("X_Co57") + # set missing radioabundance columns to zero + for col in standardcols: + if col not in dfmodel.columns and col.startswith("X_"): + dfmodel[col] = 0.0 + dfmodel["inputcellid"] = dfmodel["inputcellid"].astype(int) customcols = [col for col in dfmodel.columns if col not in standardcols] customcols.sort( key=lambda col: at.get_z_a_nucname(col) if col.startswith("X_") else (float("inf"), 0) ) # sort columns by atomic number, mass number - # set missing radioabundance columns to zero - for col in standardcols: - if col not in dfmodel.columns and col.startswith("X_"): - dfmodel[col] = 0.0 - assert modelpath is not None or filename is not None if filename is None: filename = "model.txt" - if modelpath is not None: - modelfilepath = Path(modelpath, filename) - else: - modelfilepath = Path(filename) + modelfilepath = Path(modelpath, filename) if modelpath is not None else Path(filename) with open(modelfilepath, "w", encoding="utf-8") as fmodel: if headercommentlines is not None: @@ -666,32 +746,28 @@ def save_modeldata( abundcols = [*[col for col in standardcols if col.startswith("X_")], *customcols] - # for cell in dfmodel.itertuples(): - # if dimensions == 1: - # fmodel.write(f'{cell.inputcellid:6d} {cell.velocity_outer:9.2f} {cell.logrho:10.8f} ') - # elif dimensions == 3: - # fmodel.write(f"{cell.inputcellid:6d} {cell.posx} {cell.posy} {cell.posz} {cell.rho}\n") - # - # fmodel.write(" ".join([f'{getattr(cell, col)}' for col in abundcols])) - # - # fmodel.write('\n') if dimensions == 1: for cell in dfmodel.itertuples(index=False): - fmodel.write(f"{cell.inputcellid:6d} {cell.velocity_outer:9.2f} {cell.logrho:10.8f} ") + fmodel.write(f"{cell.inputcellid:d} {cell.velocity_outer:9.2f} {cell.logrho:10.8f} ") fmodel.write(" ".join([f"{getattr(cell, col)}" for col in abundcols])) fmodel.write("\n") elif dimensions == 3: - zeroabund = " ".join(["0" for _ in abundcols]) - + # dfmodel.to_csv(fmodel, sep=" ", float_format=lambda x: "%.4e" if x > 0 else "0.0") + zeroabund = " ".join(["0.0" for _ in abundcols]) + line_end = "\n" if twolinespercell else " " for inputcellid, posxmin, posymin, poszmin, rho, *othercolvals in dfmodel[ ["inputcellid", "pos_x_min", "pos_y_min", "pos_z_min", "rho", *abundcols] ].itertuples(index=False, name=None): - fmodel.write(f"{inputcellid:6d} {posxmin} {posymin} {poszmin} {rho}\n") + fmodel.write(f"{inputcellid:d} {posxmin} {posymin} {poszmin} {rho}{line_end}") fmodel.write( " ".join( [ - f"{colvalue:.4e}" if isinstance(colvalue, float) else f"{colvalue}" + ( + (f"{colvalue:{float_format}}" if colvalue > 0.0 else "0.0") + if isinstance(colvalue, float) + else f"{colvalue}" + ) for colvalue in othercolvals ] ) @@ -711,8 +787,8 @@ def get_mgi_of_velocity_kms(modelpath: Path, velocity: float, mgilist=None) -> U velocity = float(velocity) if not mgilist: - mgilist = [mgi for mgi in modeldata.index] - arr_vouter = modeldata["velocity_outer"].values + mgilist = list(modeldata.index) + arr_vouter = modeldata["velocity_outer"].to_numpy() else: arr_vouter = np.array([modeldata["velocity_outer"][mgi] for mgi in mgilist]) @@ -720,36 +796,75 @@ def get_mgi_of_velocity_kms(modelpath: Path, velocity: float, mgilist=None) -> U if velocity < arr_vouter[index_closestvouter] or index_closestvouter + 1 >= len(mgilist): return mgilist[index_closestvouter] - elif velocity < arr_vouter[index_closestvouter + 1]: + if velocity < arr_vouter[index_closestvouter + 1]: return mgilist[index_closestvouter + 1] - elif np.isnan(velocity): + if np.isnan(velocity): return float("nan") - else: - print(f"Can't find cell with velocity of {velocity}. Velocity list: {arr_vouter}") - assert False + + print(f"Can't find cell with velocity of {velocity}. Velocity list: {arr_vouter}") + raise AssertionError @lru_cache(maxsize=8) -def get_initialabundances(modelpath: Path) -> pd.DataFrame: - """Return a list of mass fractions.""" - abundancefilepath = at.firstexisting("abundances.txt", path=modelpath, tryzipped=True) +def get_initelemabundances( + modelpath: Path = Path(), + printwarningsonly: bool = False, + dtype_backend: Literal["pyarrow", "numpy_nullable"] = "numpy_nullable", +) -> pd.DataFrame: + """Return a table of elemental mass fractions by cell from abundances.""" + abundancefilepath = at.firstexisting("abundances.txt", folder=modelpath, tryzipped=True) + + filenameparquet = at.stripallsuffixes(Path(abundancefilepath)).with_suffix(".txt.parquet") + if filenameparquet.exists() and Path(abundancefilepath).stat().st_mtime > filenameparquet.stat().st_mtime: + print(f"{abundancefilepath} has been modified after {filenameparquet}. Deleting out of date parquet file.") + filenameparquet.unlink() + + if filenameparquet.is_file(): + if not printwarningsonly: + print(f"Reading {filenameparquet}") + + abundancedata = pd.read_parquet(filenameparquet, dtype_backend=dtype_backend) + else: + if not printwarningsonly: + print(f"Reading {abundancefilepath}") + ncols = len( + pd.read_csv(at.zopen(abundancefilepath), delim_whitespace=True, header=None, comment="#", nrows=1).columns + ) + colnames = ["inputcellid", *["X_" + at.get_elsymbol(x) for x in range(1, ncols)]] + dtypes = ( + {col: "float32[pyarrow]" if col.startswith("X_") else "int32[pyarrow]" for col in colnames} + if dtype_backend == "pyarrow" + else {col: "float32" if col.startswith("X_") else "int32" for col in colnames} + ) + + abundancedata = pd.read_csv( + at.zopen(abundancefilepath), + delim_whitespace=True, + header=None, + comment="#", + names=colnames, + dtype=dtypes, + dtype_backend=dtype_backend, + ) + + if len(abundancedata) > 1000: + print(f"Saving {filenameparquet}") + abundancedata.to_parquet(filenameparquet, compression="zstd") + print(" Done.") - abundancedata = pd.read_csv(abundancefilepath, delim_whitespace=True, header=None, comment="#") abundancedata.index.name = "modelgridindex" - abundancedata.columns = ["inputcellid", *["X_" + at.get_elsymbol(x) for x in range(1, len(abundancedata.columns))]] - if len(abundancedata) > 100000: - print("abundancedata memory usage:") - abundancedata.info(verbose=False, memory_usage="deep") + if dtype_backend == "pyarrow": + abundancedata.index = abundancedata.index.astype("int32[pyarrow]") return abundancedata -def save_initialabundances( +def save_initelemabundances( dfelabundances: pd.DataFrame, abundancefilename: Union[Path, str], headercommentlines: Optional[Sequence[str]] = None, ) -> None: - """Save a DataFrame (same format as get_initialabundances) to abundances.txt. + """Save a DataFrame (same format as get_initelemabundances) to abundances.txt. columns must be: - inputcellid: integer index to match model.txt (starting from 1) - X_El: mass fraction of element with two-letter code 'El' (e.g., X_H, X_He, H_Li, ...) @@ -773,13 +888,13 @@ def save_initialabundances( fabund.write("\n".join([f"# {line}" for line in headercommentlines]) + "\n") for row in dfelabundances.itertuples(index=False): fabund.write(f" {row.inputcellid:6d} ") - fabund.write(" ".join([f"{getattr(row, colname, 0.)}" for colname in elcolnames])) + fabund.write(" ".join([f"{getattr(row, colname, 0.):.6e}" for colname in elcolnames])) fabund.write("\n") print(f"Saved {abundancefilename} (took {time.perf_counter() - timestart:.1f} seconds)") -def save_empty_abundance_file(ngrid: int, outputfilepath=Path()): +def save_empty_abundance_file(ngrid: int, outputfilepath=Path()) -> None: """Dummy abundance file with only zeros""" if Path(outputfilepath).is_dir(): outputfilepath = Path(outputfilepath) / "abundances.txt" @@ -793,7 +908,7 @@ def save_empty_abundance_file(ngrid: int, outputfilepath=Path()): # abundancedata['Z=28'] = np.ones(ngrid) dfabundances = pd.DataFrame(data=abundancedata).round(decimals=5) - dfabundances.to_csv(outputfilepath, header=False, sep="\t", index=False) + dfabundances.to_csv(outputfilepath, header=False, sep=" ", index=False) def get_dfmodel_dimensions(dfmodel: pd.DataFrame) -> int: @@ -823,7 +938,7 @@ def sphericalaverage( # dfmodel = dfmodel.query('rho > 0.').copy() dfmodel = dfmodel.copy() - celldensity = {cellindex: rho for cellindex, rho in dfmodel[["inputcellid", "rho"]].itertuples(index=False)} + celldensity = dict(dfmodel[["inputcellid", "rho"]].itertuples(index=False)) dfmodel = add_derived_cols_to_modeldata( dfmodel, ["velocity"], dimensions=3, t_model_init_seconds=t_model_init_seconds, wid_init=wid_init @@ -845,7 +960,7 @@ def sphericalaverage( highest_active_radialcellid = -1 for radialcellid, (velocity_inner, velocity_outer) in enumerate(zip(velocity_bins[:-1], velocity_bins[1:]), 1): assert velocity_outer > velocity_inner - matchedcells = dfmodel.query("vel_mid_radial > @velocity_inner and vel_mid_radial <= @velocity_outer") + matchedcells = dfmodel.query("vel_r_mid > @velocity_inner and vel_r_mid <= @velocity_outer") matchedcellrhosum = matchedcells.rho.sum() # cellidmap_3d_to_1d.update({cellid_3d: radialcellid for cellid_3d in matchedcells.inputcellid}) @@ -860,7 +975,7 @@ def sphericalaverage( # print(radialcellid, volumecorrection) if rhomean > 0.0 and dfgridcontributions is not None: - dfcellcont = dfgridcontributions.query("cellindex in @matchedcells.inputcellid.values") + dfcellcont = dfgridcontributions.query("cellindex in @matchedcells.inputcellid.to_numpy()") for particleid, dfparticlecontribs in dfcellcont.groupby("particleid"): frac_of_cellmass_avg = ( @@ -905,26 +1020,21 @@ def sphericalaverage( for column in matchedcells.columns: if column.startswith("X_") or column in ["cellYe", "q"]: - if rhomean > 0.0: - massfrac = np.dot(matchedcells[column], matchedcells.rho) / matchedcellrhosum - else: - massfrac = 0.0 + massfrac = np.dot(matchedcells[column], matchedcells.rho) / matchedcellrhosum if rhomean > 0.0 else 0.0 dictcell[column] = massfrac outcells.append(dictcell) if dfelabundances is not None: - if rhomean > 0.0: - abund_matchedcells = dfelabundances.loc[matchedcells.index] - else: - abund_matchedcells = None + abund_matchedcells = dfelabundances.loc[matchedcells.index] if rhomean > 0.0 else None dictcellabundances = {"inputcellid": radialcellid} for column in dfelabundances.columns: if column.startswith("X_"): - if rhomean > 0.0: - massfrac = np.dot(abund_matchedcells[column], matchedcells.rho) / matchedcellrhosum - else: - massfrac = 0.0 + massfrac = ( + np.dot(abund_matchedcells[column], matchedcells.rho) / matchedcellrhosum + if rhomean > 0.0 + else 0.0 + ) dictcellabundances[column] = massfrac outcellabundances.append(dictcellabundances) @@ -937,3 +1047,39 @@ def sphericalaverage( print(f" took {time.perf_counter() - timestart:.1f} seconds") return dfmodel1d, dfabundances1d, dfgridcontributions1d + + +def scale_model_to_time( + dfmodel: pd.DataFrame, + targetmodeltime_days: float, + t_model_days: Optional[float] = None, + modelmeta: Optional[dict[str, Any]] = None, +) -> tuple[pd.DataFrame, Optional[dict[str, Any]]]: + """Homologously expand model to targetmodeltime_days, reducing density and adjusting position columns to match""" + + if t_model_days is None: + assert modelmeta is not None + t_model_days = modelmeta["t_model_days"] + + timefactor = targetmodeltime_days / t_model_days + + print( + f"Adjusting t_model to {targetmodeltime_days} days (factor {timefactor}) " + "using homologous expansion of positions and densities" + ) + + for col in dfmodel.columns: + if col.startswith("pos_"): + dfmodel[col] *= timefactor + elif col == "rho": + dfmodel["rho"] *= timefactor**-3 + elif col == "logrho": + dfmodel["logrho"] += math.log10(timefactor**-3) + + if modelmeta is not None: + modelmeta["t_model_days"] = targetmodeltime_days + modelmeta.get("headercommentlines", []).append( + "scaled from {t_model_days} to {targetmodeltime_days} (no abund change from decays)" + ) + + return dfmodel, modelmeta diff --git a/artistools/inputmodel/lapuente.py b/artistools/inputmodel/lapuente.py index 257f549c6..38a3d0a6d 100755 --- a/artistools/inputmodel/lapuente.py +++ b/artistools/inputmodel/lapuente.py @@ -8,9 +8,6 @@ import artistools as at -# import os.path -# import numpy as np - def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-inputpath", "-i", default=".", help="Path of input file") @@ -34,9 +31,9 @@ def main(args=None, argsraw=None, **kwargs) -> None: datain = datain_structure.join(datain_abund, rsuffix="abund") if "Fe52" not in datain.columns: - datain.eval("Fe52 = 0.", inplace=True) + datain = datain.eval("Fe52 = 0.") if "Cr48" not in datain.columns: - datain.eval("Cr48 = 0.", inplace=True) + datain = datain.eval("Cr48 = 0.") dfmodel = pd.DataFrame( columns=[ @@ -100,7 +97,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: print(f'M_Ni56 = {tot_ni56mass / u.solMass.to("g"):.3f} solMass') at.save_modeldata(dfmodel, t_model_init_days, Path(args.outputpath, "model.txt")) - at.inputmodel.save_initialabundances(dfelabundances, Path(args.outputpath, "abundances.txt")) + at.inputmodel.save_initelemabundances(dfelabundances, Path(args.outputpath, "abundances.txt")) if __name__ == "__main__": diff --git a/artistools/inputmodel/makeartismodel.py b/artistools/inputmodel/makeartismodel.py old mode 100644 new mode 100755 index 7dcc7e3ef..60a10dbfc --- a/artistools/inputmodel/makeartismodel.py +++ b/artistools/inputmodel/makeartismodel.py @@ -6,10 +6,6 @@ import argcomplete import artistools as at -import artistools.inputmodel.downscale3dgrid -import artistools.inputmodel.energyinputfiles -import artistools.inputmodel.modelfromhydro -import artistools.inputmodel.opacityinputfile def addargs(parser: argparse.ArgumentParser) -> None: diff --git a/artistools/inputmodel/maketardismodelfromartis.py b/artistools/inputmodel/maketardismodelfromartis.py old mode 100644 new mode 100755 index 86662b25e..15d5c690d --- a/artistools/inputmodel/maketardismodelfromartis.py +++ b/artistools/inputmodel/maketardismodelfromartis.py @@ -45,10 +45,10 @@ def main(args=None, argsraw=None, **kwargs) -> None: modelpath = Path(args.inputpath) dfmodel, t_model_init_days, _ = at.inputmodel.get_modeldata_tuple( - modelpath, get_elemabundances=(args.abundtype == "elemental"), dimensions=1 + modelpath, get_elemabundances=(args.abundtype == "elemental") ) - dfmodel.eval("rho = 10 ** logrho", inplace=True) + dfmodel = dfmodel.eval("rho = 10 ** logrho") if args.abundtype == "nuclear": # nuclide abundances diff --git a/artistools/inputmodel/map_1d_to_3d_grid.py b/artistools/inputmodel/map_1d_to_3d_grid.py index 2164c22a1..045c76e0c 100644 --- a/artistools/inputmodel/map_1d_to_3d_grid.py +++ b/artistools/inputmodel/map_1d_to_3d_grid.py @@ -2,8 +2,6 @@ import artistools as at -# import pandas as pd - CLIGHT = 2.99792458e10 @@ -18,17 +16,6 @@ def change_cell_positions_to_new_time(dfgriddata, t_model_1d): return dfgriddata, wid_init -def get_cell_midpoints(dfgriddata, wid_init): - dfgriddata["posx_mid"] = dfgriddata["pos_x_min"] + (0.5 * wid_init) - dfgriddata["posy_mid"] = dfgriddata["pos_y_min"] + (0.5 * wid_init) - dfgriddata["posz_mid"] = dfgriddata["pos_z_min"] + (0.5 * wid_init) - - dfgriddata["posx_max"] = dfgriddata["pos_x_min"] + (wid_init) - dfgriddata["posy_max"] = dfgriddata["pos_y_min"] + (wid_init) - dfgriddata["posz_max"] = dfgriddata["pos_z_min"] + (wid_init) - return dfgriddata - - def map_1d_to_3d(dfgriddata, vmax, n_3d_gridcells, data_1d, t_model_1d, wid_init): modelgridindex = np.zeros(n_3d_gridcells) modelgrid_rho_3d = np.zeros(n_3d_gridcells) @@ -64,9 +51,5 @@ def map_1d_to_3d(dfgriddata, vmax, n_3d_gridcells, data_1d, t_model_1d, wid_init print(sum(modelgrid_rho_3d * (wid_init**3)) / CLIGHT) at.inputmodel.save_modeldata( - modelpath=".", - dfmodel=dfgriddata, - t_model_init_days=t_model_1d / (24 * 60 * 60), - dimensions=3, - vmax=vmax * CLIGHT, + modelpath=".", dfmodel=dfgriddata, t_model_init_days=t_model_1d / (24 * 60 * 60), vmax=vmax * CLIGHT ) diff --git a/artistools/inputmodel/maptogrid.py b/artistools/inputmodel/maptogrid.py index c3953d6f3..c5ec95190 100755 --- a/artistools/inputmodel/maptogrid.py +++ b/artistools/inputmodel/maptogrid.py @@ -16,9 +16,6 @@ itable = 40000 # wie fein Kernelfkt interpoliert wird itab = itable + 5 - -wij = np.zeros(itab + 1) - # # --maximum interaction length and step size # @@ -27,51 +24,55 @@ i1 = int(1.0 // dvtable) -igphi = 0 -# -# --normalisation constant -# -cnormk = 1.0 / math.pi -# --build tables -# -# a) v less than 1 -# -if igphi == 1: - for i in range(1, i1 + 1): - v2 = i * dvtable - v = math.sqrt(v2) - v3 = v * v2 - v4 = v * v3 - sum = 1.0 - 1.5 * v2 + 0.75 * v3 - wij[i] = cnormk * sum -else: - for i in range(1, i1 + 1): - v2 = i * dvtable - v = math.sqrt(v2) - v3 = v * v2 - sum = 1.0 - 1.5 * v2 + 0.75 * v3 - wij[i] = cnormk * sum - -# -# b) v greater than 1 -# -if igphi == 1: - for i in range(i1 + 1, itable + 1): - v2 = i * dvtable - v = math.sqrt(v2) - dif2 = 2.0 - v - sum = 0.25 * dif2 * dif2 * dif2 - wij[i] = cnormk * sum -else: - for i in range(i1 + 1, itable + 1): - v2 = i * dvtable - v = math.sqrt(v2) - dif2 = 2.0 - v - sum = 0.25 * dif2 * dif2 * dif2 - wij[i] = cnormk * sum - - -def kernelvals2(rij2: float, hmean: float) -> float: # ist schnell berechnet aber keine Gradienten +def get_wij() -> np.ndarray: + igphi = 0 + # + # --normalisation constant + # + cnormk = 1.0 / math.pi + # --build tables + # + # a) v less than 1 + # + wij = np.zeros(itab + 1) + if igphi == 1: + for i in range(1, i1 + 1): + v2 = i * dvtable + v = math.sqrt(v2) + v3 = v * v2 + v4 = v * v3 + vsum = 1.0 - 1.5 * v2 + 0.75 * v3 + wij[i] = cnormk * vsum + else: + for i in range(1, i1 + 1): + v2 = i * dvtable + v = math.sqrt(v2) + v3 = v * v2 + vsum = 1.0 - 1.5 * v2 + 0.75 * v3 + wij[i] = cnormk * vsum + + # + # b) v greater than 1 + # + if igphi == 1: + for i in range(i1 + 1, itable + 1): + v2 = i * dvtable + v = math.sqrt(v2) + dif2 = 2.0 - v + vsum = 0.25 * dif2 * dif2 * dif2 + wij[i] = cnormk * vsum + else: + for i in range(i1 + 1, itable + 1): + v2 = i * dvtable + v = math.sqrt(v2) + dif2 = 2.0 - v + vsum = 0.25 * dif2 * dif2 * dif2 + wij[i] = cnormk * vsum + + return wij + + +def kernelvals2(rij2: float, hmean: float, wij: np.ndarray) -> float: # ist schnell berechnet aber keine Gradienten hmean21 = 1.0 / (hmean * hmean) hmean31 = hmean21 / hmean v2 = rij2 * hmean21 @@ -83,11 +84,24 @@ def kernelvals2(rij2: float, hmean: float) -> float: # ist schnell berechnet ab return wtij -def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoordgrid: int = 50) -> None: +def maptogrid( + ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoordgrid: int = 50, downsamplefactor: int = 1 +) -> None: if not ejectasnapshotpath.is_file(): print(f"{ejectasnapshotpath} not found") return + # save the printed output to a log file + logfilepath = Path("maptogridlog.txt") + logfilepath.unlink(missing_ok=True) + + def logprint(*args, **kwargs): + print(*args, **kwargs) + with open(logfilepath, "at", encoding="utf-8") as logfile: + logfile.write(" ".join([str(x) for x in args]) + "\n") + + wij = get_wij() + assert ncoordgrid % 2 == 0 snapshot_columns = [ @@ -127,17 +141,21 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo snapshot_columns_used = ["id", "h", "x", "y", "z", "vx", "vy", "vz", "pmass", "rho", "p", "rho_rst", "ye"] dfsnapshot = pd.read_csv( - ejectasnapshotpath, names=snapshot_columns, delim_whitespace=True, usecols=snapshot_columns_used + ejectasnapshotpath, + names=snapshot_columns, + delim_whitespace=True, + usecols=snapshot_columns_used, + dtype_backend="pyarrow", ) - print(dfsnapshot) + if downsamplefactor > 1: + dfsnapshot = dfsnapshot.sample(len(dfsnapshot) // downsamplefactor) + + logprint(dfsnapshot) assert len(dfsnapshot.columns) == len(snapshot_columns_used) npart = len(dfsnapshot) - print("number of particles", npart) - - fpartanalysis = open(Path(outputfolderpath, "ejectapartanalysis.dat"), mode="w", encoding="utf-8") totmass = 0.0 rmax = 0.0 @@ -151,67 +169,68 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo dtextra = dtextra_seconds / 4.926e-6 # convert to geom units. - particleid = dfsnapshot.id.values - x = dfsnapshot.x.values - y = dfsnapshot.y.values - z = dfsnapshot.z.values - h = dfsnapshot.h.values - vx = dfsnapshot.vx.values - vy = dfsnapshot.vy.values - vz = dfsnapshot.vz.values - pmass = dfsnapshot.pmass.values - rho_rst = dfsnapshot.rho_rst.values - rho = dfsnapshot.rho.values - Ye = dfsnapshot.ye.values + particleid = dfsnapshot.id.to_numpy() + x = dfsnapshot["x"].to_numpy().copy() + y = dfsnapshot["y"].to_numpy().copy() + z = dfsnapshot["z"].to_numpy().copy() + h = dfsnapshot["h"].to_numpy().copy() + vx = dfsnapshot["vx"].to_numpy() + vy = dfsnapshot["vy"].to_numpy() + vz = dfsnapshot["vz"].to_numpy() + pmass = dfsnapshot["pmass"].to_numpy() + rho_rst = dfsnapshot["rho_rst"].to_numpy() + rho = dfsnapshot["rho"].to_numpy() + Ye = dfsnapshot["ye"].to_numpy() - for n in range(npart): - totmass = totmass + pmass[n] + with open(Path(outputfolderpath, "ejectapartanalysis.dat"), mode="w", encoding="utf-8") as fpartanalysis: + for n in range(npart): + totmass = totmass + pmass[n] - dis = math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2) # original dis + dis = math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2) # original dis - x[n] += vx[n] * dtextra - y[n] += vy[n] * dtextra - z[n] += vz[n] * dtextra + x[n] += vx[n] * dtextra + y[n] += vy[n] * dtextra + z[n] += vz[n] * dtextra - # actually we should also extrapolate smoothing length h unless we disrgard it below + # actually we should also extrapolate smoothing length h unless we disrgard it below - # extrapolate h such that ratio betwen dis and h remains constant - h[n] = h[n] / dis * math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2) + # extrapolate h such that ratio between dis and h remains constant + h[n] = h[n] / dis * math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2) - dis = math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2) # possibly new distance + dis = math.sqrt(x[n] ** 2 + y[n] ** 2 + z[n] ** 2) # possibly new distance - rmean = rmean + dis + rmean = rmean + dis - rmax = max(rmax, dis) + rmax = max(rmax, dis) - hmean = hmean + h[n] + hmean = hmean + h[n] - hmin = min(hmin, h[n]) + hmin = min(hmin, h[n]) - vtot = math.sqrt(vx[n] ** 2 + vy[n] ** 2 + vz[n] ** 2) + vtot = math.sqrt(vx[n] ** 2 + vy[n] ** 2 + vz[n] ** 2) - vrad = (vx[n] * x[n] + vy[n] * y[n] + vz[n] * z[n]) / dis # radial velocity + vrad = (vx[n] * x[n] + vy[n] * y[n] + vz[n] * z[n]) / dis # radial velocity - if vtot > vrad: - vperp = math.sqrt(vtot * vtot - vrad * vrad) # velicty perpendicular - else: - vperp = 0.0 # if we rxtrapolate roundoff error can lead to Nan, ly? + # velocity perpendicular + # if we extrapolate roundoff error can lead to Nan, ly? + vperp = math.sqrt(vtot * vtot - vrad * vrad) if vtot > vrad else 0.0 - vratiomean = vratiomean + vperp / vrad + vratiomean = vratiomean + vperp / vrad - # output some ejecta properties in file + # output some ejecta properties in file - fpartanalysis.write(f"{dis} {h[n]} {h[n] / dis} {vrad} {vperp} {vtot}\n") + fpartanalysis.write(f"{dis} {h[n]} {h[n] / dis} {vrad} {vperp} {vtot}\n") + logprint("saved ejectapartanalysis.dat") rmean = rmean / npart hmean = hmean / npart vratiomean = vratiomean / npart - print("total mass of sph particle, max, mean distance", totmass, rmax, rmean) - print("smoothing length min, mean", hmin, hmean) - print("ratio between vrad and vperp mean", vratiomean) + logprint(f"total mass of sph particle {totmass} max dist {rmax} mean dist {rmean}") + logprint(f"smoothing length min {hmin} mean {hmean}") + logprint("ratio between vrad and vperp mean", vratiomean) # check maybe cm and correct by shifting @@ -230,28 +249,33 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo dy = dx dz = dx - grho = np.zeros((ncoordgrid + 1, ncoordgrid + 1, ncoordgrid + 1)) - gye = np.zeros((ncoordgrid + 1, ncoordgrid + 1, ncoordgrid + 1)) - gparticlecounter = np.zeros((ncoordgrid + 1, ncoordgrid + 1, ncoordgrid + 1), dtype=int) + grho = np.zeros((ncoordgrid, ncoordgrid, ncoordgrid)) + gye = np.zeros((ncoordgrid, ncoordgrid, ncoordgrid)) + gparticlecounter = np.zeros((ncoordgrid, ncoordgrid, ncoordgrid), dtype=int) particle_rho_contribs = {} - print("grid properties", x0, dx, x0 + dx * (ncoordgrid - 1)) + logprint(f"grid properties {x0=}, {dx=}, {x0 + dx * (ncoordgrid - 1)=}") - arrgx = x0 + dx * (np.arange(ncoordgrid + 1) - 1) + arrgx = x0 + dx * np.arange(ncoordgrid) arrgy = arrgx arrgz = arrgx + particlesused = set() + particlesinsidegrid = set() + for n in range(npart): maxdist = 2.0 * h[n] maxdist2 = maxdist**2 - ilow = max(math.floor((x[n] - maxdist - x0) / dx), 1) - ihigh = min(math.ceil((x[n] + maxdist - x0) / dx), ncoordgrid) - jlow = max(math.floor((y[n] - maxdist - y0) / dy), 1) - jhigh = min(math.ceil((y[n] + maxdist - y0) / dy), ncoordgrid) - klow = max(math.floor((z[n] - maxdist - z0) / dz), 1) - khigh = min(math.ceil((z[n] + maxdist - z0) / dz), ncoordgrid) + ilow = max(math.floor((x[n] - maxdist - x0) / dx), 0) + ihigh = min(math.ceil((x[n] + maxdist - x0) / dx), ncoordgrid - 1) + jlow = max(math.floor((y[n] - maxdist - y0) / dy), 0) + jhigh = min(math.ceil((y[n] + maxdist - y0) / dy), ncoordgrid - 1) + klow = max(math.floor((z[n] - maxdist - z0) / dz), 0) + khigh = min(math.ceil((z[n] + maxdist - z0) / dz), ncoordgrid - 1) + if min(ihigh, jhigh, khigh) >= 1 and max(ilow, jlow, klow) <= ncoordgrid: + particlesinsidegrid.add(n) # check some min max # ... kernel reweighting ? @@ -290,7 +314,7 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo # or via neighbors - not yet implemented if dis2 <= maxdist2: - wtij = kernelvals2(dis2, h[n]) + wtij = kernelvals2(dis2, h[n], wij) # USED PREVIOUSLY: less accurate? # grho_contrib = pmass[n] * wtij @@ -307,6 +331,21 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo # count number of particles contributing to each grid cell gparticlecounter[i, j, k] += 1 + particlesused.add(n) + + logprint( + f"particles with any cell contribution: {len(particlesused)} of {len(particlesinsidegrid)} inside grid out of" + f" {npart} total" + ) + unusedparticles = [n for n in range(npart) if n not in particlesused] + for n in unusedparticles: + loc_i = math.floor((x[n] - x0) / dx) + loc_j = math.floor((y[n] - y0) / dy) + loc_k = math.floor((z[n] - z0) / dz) + # ignore particles outside grid boundary + if min(loc_i, loc_j, loc_k) < 0 or max(loc_i, loc_j, loc_k) > ncoordgrid - 1: + continue + logprint(f"particle {n} is totally unused but located in cell {loc_i} {loc_j} {loc_k}") with np.errstate(divide="ignore", invalid="ignore"): gye = np.divide(gye, grho) @@ -314,8 +353,9 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo with open(Path(outputfolderpath, "gridcontributions.txt"), "w", encoding="utf-8") as fcontribs: fcontribs.write("particleid cellindex frac_of_cellmass\n") for (n, i, j, k), rho_contrib in particle_rho_contribs.items(): - gridindex = ((k - 1) * ncoordgrid + (j - 1)) * ncoordgrid + (i - 1) + 1 + gridindex = (k * ncoordgrid + j) * ncoordgrid + i + 1 fcontribs.write(f"{particleid[n]} {gridindex} {rho_contrib / grho[i, j, k]}\n") + logprint("saved gridcontributions.txt") # check some stuff on the grid @@ -324,39 +364,38 @@ def maptogrid(ejectasnapshotpath: Path, outputfolderpath: Union[Path, str], ncoo gmass = np.sum(grho) * dx * dy * dz # nzero = np.count_nonzero(grho[1:][1:][1:] < 1.e-20) - for i in range(1, ncoordgrid + 1): - gx = x0 + dx * (i - 1) - for j in range(1, ncoordgrid + 1): - gy = y0 + dy * (j - 1) - for k in range(1, ncoordgrid + 1): + for i in range(ncoordgrid): + gx = x0 + dx * i + for j in range(ncoordgrid): + gy = y0 + dy * j + for k in range(1, ncoordgrid): # how many cells with rho=0? if grho[i, j, k] < 1.0e-20: nzero = nzero + 1 - gz = z0 + dz * (k - 1) + gz = z0 + dz * k dis = math.sqrt(gx * gx + gy * gy + gz * gz) if grho[i, j, k] < 1.0e-20 and dis < rmean: nzerocentral = nzerocentral + 1 - print(f"mass on grid from rho*V: {gmass} particles: {totmass}") + logprint( + f"{'WARNING!' if gmass / totmass < 0.9 else ''} mass on grid from rho*V: {gmass} mass of particles: {totmass} " + ) - print( - "number of cells with zero rho, total num of cells, fraction of cells w rho=0", - nzero, - ncoordgrid**3, - (nzero) / (ncoordgrid**3), + logprint( + f"number of cells with rho=0 {nzero}, total num of cells {ncoordgrid**3}, fraction of cells with rho=0:" + f" {(nzero) / (ncoordgrid**3)}" ) - print( - "number of central cells (dis None: parser.add_argument("-inputpath", "-i", default=".", help="Path to ejectasnapshot") parser.add_argument( "-ncoordgrid", type=int, default=50, help="Number of grid positions per axis (numcells = ncoordgrid^3)" ) + parser.add_argument( + "-downsamplefactor", + type=int, + default=1, + help="Randomly sample particles, reducing the number by this factor (e.g. 2 will ignore half of the particles)", + ) parser.add_argument("-outputfolderpath", "-o", default=".", help="Path for output files") @@ -401,7 +448,12 @@ def main(args=None, argsraw=None, **kwargs) -> None: ejectasnapshotpath = Path(args.inputpath, "ejectasnapshot.dat") - maptogrid(ejectasnapshotpath=ejectasnapshotpath, ncoordgrid=args.ncoordgrid, outputfolderpath=args.outputfolderpath) + maptogrid( + ejectasnapshotpath=ejectasnapshotpath, + ncoordgrid=args.ncoordgrid, + outputfolderpath=args.outputfolderpath, + downsamplefactor=args.downsamplefactor, + ) if __name__ == "__main__": diff --git a/artistools/inputmodel/modelfromhydro.py b/artistools/inputmodel/modelfromhydro.py old mode 100644 new mode 100755 index 08bc40842..5b9d95bfc --- a/artistools/inputmodel/modelfromhydro.py +++ b/artistools/inputmodel/modelfromhydro.py @@ -3,7 +3,7 @@ import argparse import datetime import math -import os.path +import sys from pathlib import Path import argcomplete @@ -77,9 +77,9 @@ def get_merger_time_geomunits(pathtogriddata: Path) -> float: mergertime_geomunits = float(fmergertimefile.readline()) print(f"Found simulation merger time to be {mergertime_geomunits} ({mergertime_geomunits * 4.926e-6} s) ") return mergertime_geomunits - else: - print('Make file "tmerger.txt" with time of merger in geom units') - quit() + + print('Make file "tmerger.txt" with time of merger in geom units') + sys.exit(1) def get_snapshot_time_geomunits(pathtogriddata: Path) -> tuple[float, float]: @@ -88,11 +88,11 @@ def get_snapshot_time_geomunits(pathtogriddata: Path) -> tuple[float, float]: snapshotinfofiles = glob.glob(str(Path(pathtogriddata) / "*_info.dat*")) if not snapshotinfofiles: print("No info file found for dumpstep") - quit() + sys.exit(1) if len(snapshotinfofiles) > 1: print("Too many sfho_info.dat files found") - quit() + sys.exit(1) snapshotinfofile = Path(snapshotinfofiles[0]) if snapshotinfofile.is_file(): @@ -109,23 +109,11 @@ def get_snapshot_time_geomunits(pathtogriddata: Path) -> tuple[float, float]: else: print("Could not find snapshot info file to get simulation time") - quit() + sys.exit(1) return simulation_end_time_geomunits, mergertime_geomunits -def scale_model_to_time(targetmodeltime_days, t_model_days, modeldata): - timefactor = targetmodeltime_days / t_model_days - modeldata.eval("pos_x_min = pos_x_min * @timefactor", inplace=True) - modeldata.eval("pos_y_min = pos_y_min * @timefactor", inplace=True) - modeldata.eval("pos_z_min = pos_z_min * @timefactor", inplace=True) - modeldata.eval("rho = rho * @timefactor ** -3", inplace=True) - print( - f"Adjusting t_model to {targetmodeltime_days} days (factor {timefactor}) " - "using homologous expansion of positions and densities" - ) - - def read_griddat_file(pathtogriddata, targetmodeltime_days=None, minparticlespercell=0): griddatfilepath = Path(pathtogriddata) / "grid.dat" @@ -133,14 +121,13 @@ def read_griddat_file(pathtogriddata, targetmodeltime_days=None, minparticlesper simulation_end_time_geomunits, mergertime_geomunits = get_snapshot_time_geomunits(pathtogriddata) griddata = pd.read_csv(griddatfilepath, delim_whitespace=True, comment="#", skiprows=3) - griddata.rename( + griddata = griddata.rename( columns={ "gridindex": "inputcellid", "pos_x": "pos_x_min", "pos_y": "pos_y_min", "pos_z": "pos_z_min", }, - inplace=True, ) # griddata in geom units griddata["rho"] = np.nan_to_num(griddata["rho"], nan=0.0) @@ -152,17 +139,17 @@ def read_griddat_file(pathtogriddata, targetmodeltime_days=None, minparticlesper factor_position = 1.478 # in km km_to_cm = 1e5 - griddata.eval("pos_x_min = pos_x_min * @factor_position * @km_to_cm", inplace=True) - griddata.eval("pos_y_min = pos_y_min * @factor_position * @km_to_cm", inplace=True) - griddata.eval("pos_z_min = pos_z_min * @factor_position * @km_to_cm", inplace=True) + griddata["pos_x_min"] = griddata["pos_x_min"] * factor_position * km_to_cm + griddata["pos_y_min"] = griddata["pos_y_min"] * factor_position * km_to_cm + griddata["pos_z_min"] = griddata["pos_z_min"] * factor_position * km_to_cm griddata["rho"] = griddata["rho"] * 6.176e17 # convert to g/cm3 - with open(griddatfilepath, "r", encoding="utf-8") as gridfile: + with open(griddatfilepath, encoding="utf-8") as gridfile: ngrid = int(gridfile.readline().split()[0]) if ngrid != len(griddata["inputcellid"]): print("length of file and ngrid don't match") - quit() + sys.exit(1) extratime_geomunits = float(gridfile.readline().split()[0]) xmax = abs(float(gridfile.readline().split()[0])) xmax = (xmax * factor_position) * (u.km).to(u.cm) @@ -183,16 +170,18 @@ def read_griddat_file(pathtogriddata, targetmodeltime_days=None, minparticlesper ) if targetmodeltime_days is not None: - scale_model_to_time(targetmodeltime_days, t_model_days, griddata) + dfmodel, modelmeta = at.inputmodel.scale_model_to_time( + targetmodeltime_days=targetmodeltime_days, t_model_days=t_model_days, dfmodel=griddata + ) t_model_days = targetmodeltime_days if minparticlespercell > 0: ncoordgridx = round(len(griddata) ** (1.0 / 3.0)) xmax = -griddata.pos_x_min.min() wid_init = 2 * xmax / ncoordgridx - filter = np.logical_and(griddata.tracercount < minparticlespercell, griddata.rho > 0.0) - n_ignored = np.count_nonzero(filter) - mass_ignored = griddata.loc[filter].rho.sum() * wid_init**3 / 1.989e33 + cellfilter = np.logical_and(griddata.tracercount < minparticlespercell, griddata.rho > 0.0) + n_ignored = np.count_nonzero(cellfilter) + mass_ignored = griddata.loc[cellfilter].rho.sum() * wid_init**3 / 1.989e33 mass_orig = griddata.rho.sum() * wid_init**3 / 1.989e33 print( @@ -211,7 +200,7 @@ def read_mattia_grid_data_file(pathtogriddata): griddatfilepath = Path(pathtogriddata) / "1D_m0.01_v0.1.txt" griddata = pd.read_csv(griddatfilepath, delim_whitespace=True, comment="#", skiprows=1) - with open(griddatfilepath, "r") as gridfile: + with open(griddatfilepath) as gridfile: t_model = float(gridfile.readline()) print(f"t_model {t_model} seconds") xmax = max(griddata["posx"]) @@ -346,10 +335,11 @@ def makemodelfromgriddata( if getattr(args, "getcellopacityfromYe", False): at.inputmodel.opacityinputfile.opacity_by_Ye(outputpath, dfmodel) - if os.path.isfile(Path(gridfolderpath, "gridcontributions.txt")): - dfgridcontributions = at.inputmodel.rprocess_from_trajectory.get_gridparticlecontributions(gridfolderpath) - else: - dfgridcontributions = None + dfgridcontributions = ( + at.inputmodel.rprocess_from_trajectory.get_gridparticlecontributions(gridfolderpath) + if Path(gridfolderpath, "gridcontributions.txt").is_file() + else None + ) if traj_root is not None: print(f"Nuclear network abundances from {traj_root} will be used") @@ -387,11 +377,11 @@ def makemodelfromgriddata( dfgridcontributions, Path(outputpath, "gridcontributions.txt") ) - headercommentlines.append(f"generated at (UTC): {datetime.datetime.utcnow()}") + headercommentlines.append(f"generated at (UTC): {datetime.datetime.now(tz=datetime.timezone.utc)}") if traj_root is not None: print(f'Writing to {Path(outputpath) / "abundances.txt"}...') - at.inputmodel.save_initialabundances( + at.inputmodel.save_initelemabundances( dfelabundances=dfelabundances, abundancefilename=outputpath, headercommentlines=headercommentlines ) else: @@ -448,16 +438,14 @@ def main(args=None, argsraw=None, **kwargs): argcomplete.autocomplete(parser) args = parser.parse_args(argsraw) + pd.options.mode.copy_on_write = True gridfolderpath = args.gridfolderpath if not Path(gridfolderpath, "grid.dat").is_file() or not Path(gridfolderpath, "gridcontributions.txt").is_file(): print("grid.dat and gridcontributions.txt are required. Run artistools-maptogrid") return # at.inputmodel.maptogrid.main() - if args.outputpath is None: - outputpath = Path(f"artismodel_{args.dimensions}d") - else: - outputpath = Path(args.outputpath) + outputpath = Path(f"artismodel_{args.dimensions}d") if args.outputpath is None else Path(args.outputpath) outputpath.mkdir(parents=True, exist_ok=True) diff --git a/artistools/inputmodel/plotdensity.py b/artistools/inputmodel/plotdensity.py index 06b9bb365..5b04a81d4 100755 --- a/artistools/inputmodel/plotdensity.py +++ b/artistools/inputmodel/plotdensity.py @@ -1,13 +1,10 @@ #!/usr/bin/env python3 # PYTHON_ARGCOMPLETE_OK import argparse -import math from pathlib import Path import argcomplete import matplotlib.pyplot as plt -import numpy as np -import pandas as pd import artistools as at diff --git a/artistools/inputmodel/recombinationenergy.py b/artistools/inputmodel/recombinationenergy.py index 5746fa1a2..9a19eab04 100755 --- a/artistools/inputmodel/recombinationenergy.py +++ b/artistools/inputmodel/recombinationenergy.py @@ -1,7 +1,6 @@ #!/usr/bin/env python3 # PYTHON_ARGCOMPLETE_OK import argparse -import math from pathlib import Path import argcomplete @@ -10,10 +9,6 @@ import pandas as pd import artistools as at -import artistools.inputmodel - -# import os.path -# import matplotlib def get_model_recombenergy(dfbinding, args): @@ -48,11 +43,9 @@ def get_model_recombenergy(dfbinding, args): elsymb = at.get_elsymbol(atomic_number) massnumber = int(species.lstrip("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")) matchrows = dfbinding.query("Z == @atomic_number") - if matchrows.empty: - binding_en_ev = 0.0 - # print(f'No binding energy for Z={atomic_number}') - else: - binding_en_ev = matchrows.iloc[0]["TotBEn"] + + binding_en_ev = 0.0 if matchrows.empty else matchrows.iloc[0]["TotBEn"] + # print(species, atomic_number, massnumber, el_binding_en_ev) contrib_binding_en_ev = speciesabund_g / (massnumber * amu_g) * binding_en_ev @@ -97,7 +90,7 @@ def get_particle_elec_binding_energy_per_gram(traj_root, dictbinding, particleid # print(dftrajnucabund) # assert frac_unaccounted < 0.3 - dftrajnucabund.eval("recombenergy_ev_per_gram = Z_be_tot_ev * massfrac / (Z + N) / @amu_g", inplace=True) + dftrajnucabund = dftrajnucabund.eval("recombenergy_ev_per_gram = Z_be_tot_ev * massfrac / (Z + N) / @amu_g") # contrib_binding_en_ev = speciesabund_g / (massnumber * amu_g) * binding_en_ev @@ -113,9 +106,9 @@ def get_particle_nucenergy_released(traj_root, particleid, tmin_s, time_s_end): traj_root=traj_root, particleid=particleid, memberfilename=memberfilename ) as fthermo: dfthermo = pd.read_csv(fthermo, delim_whitespace=True, usecols=["#count", "time/s", "Qdot", "Ye"]) - dfthermo.rename(columns={"time/s": "time_s"}, inplace=True) - dfthermo.query("time_s >= @tmin_s", inplace=True) - dfthermo.query("time_s <= @time_s_end", inplace=True) + dfthermo = dfthermo.rename(columns={"time/s": "time_s"}) + dfthermo = dfthermo.query("time_s >= @tmin_s") + dfthermo = dfthermo.query("time_s <= @time_s_end") en_released_ev_per_gram = np.trapz(y=dfthermo["Qdot"], x=dfthermo["time_s"]) * erg_to_ev # print(dfthermo) return en_released_ev_per_gram @@ -125,9 +118,9 @@ def get_particles_recomb_nuc_energy(traj_root, dfbinding): dfsnapshot = at.inputmodel.modelfromhydro.read_ejectasnapshot( "/Users/luke/Library/Mobile Documents/com~apple~CloudDocs/Archive/Astronomy/Mergers/SFHo_snapshot" ) - dfsnapshot.sort_values("ye", inplace=True) + dfsnapshot = dfsnapshot.sort_values("ye") - dictbinding = {Z: be_tot_ev for Z, be_tot_ev in dfbinding[["Z", "TotBEn"]].itertuples(index=False)} + dictbinding = dict(dfbinding[["Z", "TotBEn"]].itertuples(index=False)) tmin_s = 10 time_s = 6 * 3600 @@ -137,7 +130,7 @@ def get_particles_recomb_nuc_energy(traj_root, dfbinding): ye_list = [] elecbinding_en_list = [] nuclear_released_en_list = [] - for particleid, ye, pmass in dfsnapshot[["id", "ye", "pmass"]].itertuples(index=False): + for particleid, ye, _pmass in dfsnapshot[["id", "ye", "pmass"]].itertuples(index=False): try: elecbinding_en = get_particle_elec_binding_energy_per_gram( traj_root=traj_root, dictbinding=dictbinding, particleid=particleid, time_s=time_s @@ -213,7 +206,7 @@ def main(args=None, argsraw=None, **kwargs): argcomplete.autocomplete(parser) args = parser.parse_args(argsraw) - with open(at.get_config()["path_datadir"] / "ElBiEn_2007.txt", "r") as fbinding: + with open(at.get_config()["path_datadir"] / "ElBiEn_2007.txt") as fbinding: for _ in range(11): header = fbinding.readline().lstrip(" #").split() # print(header) diff --git a/artistools/inputmodel/rprocess_from_trajectory.py b/artistools/inputmodel/rprocess_from_trajectory.py old mode 100644 new mode 100755 index 33ec629f9..398d356f4 --- a/artistools/inputmodel/rprocess_from_trajectory.py +++ b/artistools/inputmodel/rprocess_from_trajectory.py @@ -21,7 +21,7 @@ def get_elemabund_from_nucabund(dfnucabund: pd.DataFrame) -> dict[str, float]: - """return a dictionary of elemental abundances from nuclear abundance DataFrame""" + """Return a dictionary of elemental abundances from nuclear abundance DataFrame.""" dictelemabund: dict[str, float] = {} for atomic_number in range(1, dfnucabund.Z.max() + 1): dictelemabund[f"X_{at.get_elsymbol(atomic_number)}"] = dfnucabund.query( @@ -31,8 +31,7 @@ def get_elemabund_from_nucabund(dfnucabund: pd.DataFrame) -> dict[str, float]: def open_tar_file_or_extracted(traj_root: Path, particleid: int, memberfilename: str): - """ - trajectory files are generally stored as {particleid}.tar.xz, but this is slow + """Trajectory files are generally stored as {particleid}.tar.xz, but this is slow to access, so first check for extracted files, or decompressed .tar files, which are much faster to access. @@ -47,7 +46,7 @@ def open_tar_file_or_extracted(traj_root: Path, particleid: int, memberfilename: tarfile.open(tarfilepath, "r:*").extract(path=Path(traj_root, str(particleid)), member=memberfilename) if path_extracted_file.is_file(): - return open(path_extracted_file, mode="r", encoding="utf-8") + return open(path_extracted_file, encoding="utf-8") if not tarfilepath.is_file(): print(f"No network data found for particle {particleid} (so can't access {memberfilename})") @@ -63,20 +62,21 @@ def open_tar_file_or_extracted(traj_root: Path, particleid: int, memberfilename: return io.StringIO(extractedfile.read().decode("utf-8")) print(f"Member {memberfilename} not found in {tarfilepath}") - assert False + raise AssertionError @lru_cache(maxsize=16) def get_dfevol(traj_root: Path, particleid: int) -> pd.DataFrame: with open_tar_file_or_extracted(traj_root, particleid, "./Run_rprocess/evol.dat") as evolfile: - dfevol = pd.read_table( + dfevol = pd.read_csv( evolfile, sep=r"\s+", comment="#", usecols=[0, 1], names=["nstep", "timesec"], engine="c", - dtype={0: int, 1: float}, + dtype={0: "int32[pyarrow]", 1: "float32[pyarrow]"}, + dtype_backend="pyarrow", ) return dfevol @@ -93,7 +93,7 @@ def get_closest_network_timestep( dfevol: pd.DataFrame = get_dfevol(traj_root, particleid) if cond == "nearest": - idx = np.abs(dfevol.timesec.values - timesec).argmin() + idx = np.abs(dfevol.timesec.to_numpy() - timesec).argmin() elif cond == "greaterthan": return dfevol.query("timesec > @timesec").nstep.min() @@ -102,9 +102,9 @@ def get_closest_network_timestep( return dfevol.query("timesec < @timesec").nstep.max() else: - assert False + raise AssertionError - nts: int = dfevol.nstep.values[idx] + nts: int = dfevol.nstep.to_numpy()[idx] return nts @@ -119,10 +119,10 @@ def get_trajectory_timestepfile_nuc_abund( with open_tar_file_or_extracted(traj_root, particleid, memberfilename) as trajfile: try: _, str_t_model_init_seconds, _, rho, _, _ = trajfile.readline().split() - except ValueError: + except ValueError as exc: print(f"Problem with {memberfilename}") - raise ValueError(f"Problem with {memberfilename}") - assert False + raise ValueError(f"Problem with {memberfilename}") from exc + trajfile.seek(0) t_model_init_seconds = float(str_t_model_init_seconds) @@ -133,11 +133,12 @@ def get_trajectory_timestepfile_nuc_abund( colspecs=[(0, 4), (4, 8), (8, 21)], engine="c", names=["N", "Z", "log10abund"], - dtype={0: int, 1: int, 2: float}, + dtype={0: "int32[pyarrow]", 1: "int32[pyarrow]", 2: "float32[pyarrow]"}, + dtype_backend="pyarrow", ) # in case the files are inconsistent, switch to an adaptive reader - # dfnucabund = pd.read_table( + # dfnucabund = pd.read_csv( # trajfile, # skip_blank_lines=True, # skiprows=1, @@ -149,7 +150,7 @@ def get_trajectory_timestepfile_nuc_abund( # ) # dfnucabund.eval('abund = 10 ** log10abund', inplace=True) - dfnucabund.eval("massfrac = (N + Z) * (10 ** log10abund)", inplace=True) + dfnucabund["massfrac"] = (dfnucabund["N"] + dfnucabund["Z"]) * (10 ** dfnucabund["log10abund"]) # dfnucabund.eval('A = N + Z', inplace=True) # dfnucabund.query('abund > 0.', inplace=True) @@ -162,16 +163,26 @@ def get_trajectory_timestepfile_nuc_abund( def get_trajectory_qdotintegral(particleid: int, traj_root: Path, nts_max: int, t_model_s: float) -> float: - """initial cell energy [erg/g]""" + """Initial cell energy [erg/g].""" with open_tar_file_or_extracted(traj_root, particleid, "./Run_rprocess/energy_thermo.dat") as enthermofile: - dfthermo: pd.DataFrame = pd.read_table( - enthermofile, sep=r"\s+", usecols=["time/s", "Qdot"], engine="c", dtype={0: float, 1: float} - ) - dfthermo.rename(columns={"time/s": "time_s"}, inplace=True) + try: + dfthermo: pd.DataFrame = pd.read_csv( + enthermofile, + sep=r"\s+", + usecols=["time/s", "Qdot"], + engine="c", + dtype={0: "float32[pyarrow]", 1: "float32[pyarrow]"}, + dtype_backend="pyarrow", + ) + except pd.errors.EmptyDataError: + print(f"Problem with file {enthermofile}") + raise + + dfthermo = dfthermo.rename(columns={"time/s": "time_s"}) startindex: int = int(np.argmax(dfthermo["time_s"] >= 1)) # start integrating at this number of seconds assert all(dfthermo["Qdot"][startindex : nts_max + 1] > 0.0) - dfthermo.eval("Qdot_expansionadjusted = Qdot * time_s / @t_model_s", inplace=True) + dfthermo["Qdot_expansionadjusted"] = dfthermo["Qdot"] * dfthermo["time_s"] / t_model_s qdotintegral: float = np.trapz( y=dfthermo["Qdot_expansionadjusted"][startindex : nts_max + 1], @@ -189,9 +200,8 @@ def get_trajectory_abund_q( nts: Optional[int] = None, getqdotintegral: bool = False, ) -> dict[str, float]: - """ - get the nuclear mass fractions (and Qdotintegral) for a particle particle number as a given time - nts: GSI network timestep number + """Get the nuclear mass fractions (and Qdotintegral) for a particle particle number as a given time + nts: GSI network timestep number. """ assert t_model_s is not None or nts is not None try: @@ -214,7 +224,7 @@ def get_trajectory_abund_q( return {} massfractotal = dftrajnucabund.massfrac.sum() - dftrajnucabund.query("Z >= 1", inplace=True) + dftrajnucabund = dftrajnucabund.loc[dftrajnucabund["Z"] >= 1] dftrajnucabund["nucabundcolname"] = [ f"X_{at.get_elsymbol(int(row.Z))}{int(row.N + row.Z)}" for row in dftrajnucabund.itertuples() @@ -261,9 +271,9 @@ def get_modelcellabundance( # adjust frac_of_cellmass for missing particles cell_frac_sum = sum([frac_of_cellmass for _, frac_of_cellmass in contribparticles]) - nucabundcolnames = set( - [col for particleid in dfthiscellcontribs.particleid for col in dict_traj_nuc_abund.get(particleid, {}).keys()] - ) + nucabundcolnames = { + col for particleid in dfthiscellcontribs.particleid for col in dict_traj_nuc_abund.get(particleid, {}) + } row = { nucabundcolname: sum( @@ -291,7 +301,13 @@ def get_gridparticlecontributions(gridcontribpath: Union[Path, str]) -> pd.DataF dfcontribs = pd.read_csv( Path(gridcontribpath, "gridcontributions.txt"), delim_whitespace=True, - dtype={0: int, 1: int, 2: float, 3: float}, + dtype={ + 0: "int32[pyarrow]", + 1: "int32[pyarrow]", + 2: "float32[pyarrow]", + 3: "float32[pyarrow]", + }, + dtype_backend="pyarrow", ) return dfcontribs @@ -317,7 +333,7 @@ def filtermissinggridparticlecontributions(traj_root: Path, dfcontribs: pd.DataF ) # after filtering, frac_of_cellmass_includemissing will still include particles with rho but no abundance data # frac_of_cellmass will exclude particles with no abundances - dfcontribs.eval("frac_of_cellmass_includemissing = frac_of_cellmass", inplace=True) + dfcontribs["frac_of_cellmass_includemissing"] = dfcontribs["frac_of_cellmass"] # dfcontribs.query('particleid not in @missing_particleids', inplace=True) dfcontribs.loc[dfcontribs.eval("particleid in @missing_particleids"), "frac_of_cellmass"] = 0.0 @@ -359,11 +375,11 @@ def filtermissinggridparticlecontributions(traj_root: Path, dfcontribs: pd.DataF return dfcontribs -def save_gridparticlecontributions(dfcontribs: pd.DataFrame, gridcontribpath): +def save_gridparticlecontributions(dfcontribs: pd.DataFrame, gridcontribpath) -> None: gridcontribpath = Path(gridcontribpath) if gridcontribpath.is_dir(): gridcontribpath = Path(gridcontribpath, "gridcontributions.txt") - dfcontribs.to_csv(gridcontribpath, sep=" ", index=False) + dfcontribs.to_csv(gridcontribpath, sep=" ", index=False, float_format="%.7e") def add_abundancecontributions( @@ -373,7 +389,7 @@ def add_abundancecontributions( traj_root: Path, minparticlespercell: int = 0, ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: - """contribute trajectory network calculation abundances to model cell abundances""" + """Contribute trajectory network calculation abundances to model cell abundances.""" t_model_s = t_model_days_incpremerger * 86400 dfcontribs = dfgridcontributions @@ -386,7 +402,7 @@ def add_abundancecontributions( if len(dfthiscellcontribs) >= minparticlespercell ] - dfcontribs.query("cellindex in @active_inputcellids", inplace=True) + dfcontribs = dfcontribs.query("cellindex in @active_inputcellids") dfcontribs = filtermissinggridparticlecontributions(traj_root, dfcontribs) active_inputcellids = dfcontribs.cellindex.unique() active_inputcellcount = len(active_inputcellids) @@ -444,9 +460,9 @@ def add_abundancecontributions( timestart = time.perf_counter() print("Creating dfnucabundances...", end="", flush=True) dfnucabundances = pd.DataFrame(listcellnucabundances) - dfnucabundances.set_index("inputcellid", drop=False, inplace=True) + dfnucabundances = dfnucabundances.set_index("inputcellid", drop=False) dfnucabundances.index.name = None - dfnucabundances.fillna(0.0, inplace=True) + dfnucabundances = dfnucabundances.fillna(0.0) print(f" took {time.perf_counter() - timestart:.1f} seconds") timestart = time.perf_counter() @@ -471,7 +487,7 @@ def add_abundancecontributions( f"X_{at.get_elsymbol(atomic_number)}": ( dfnucabundances.eval( f'{" + ".join(elemisotopes[atomic_number])}', - # engine="python" if len(elemisotopes[atomic_number]) > 31 else None + engine="python" if len(elemisotopes[atomic_number]) > 31 else None, ) if atomic_number in elemisotopes else np.zeros(len(dfnucabundances)) @@ -486,15 +502,15 @@ def add_abundancecontributions( dfelabundances = dfmodel[["inputcellid"]].merge( dfelabundances_partial, how="left", left_on="inputcellid", right_on="inputcellid" ) - dfnucabundances.set_index("inputcellid", drop=False, inplace=True) + dfnucabundances = dfnucabundances.set_index("inputcellid", drop=False) dfnucabundances.index.name = None - dfelabundances.fillna(0.0, inplace=True) + dfelabundances = dfelabundances.fillna(0.0) print(f" took {time.perf_counter() - timestart:.1f} seconds") print(f" there are {nuclidesincluded} nuclides from {elementsincluded} elements included") timestart = time.perf_counter() print("Merging isotopic abundances into dfmodel...", end="", flush=True) dfmodel = dfmodel.merge(dfnucabundances, how="left", left_on="inputcellid", right_on="inputcellid") - dfmodel.fillna(0.0, inplace=True) + dfmodel = dfmodel.fillna(0.0) print(f" took {time.perf_counter() - timestart:.1f} seconds") return dfmodel, dfelabundances, dfcontribs @@ -504,7 +520,7 @@ def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-outputpath", "-o", default=".", help="Path for output files") -def main(args=None, argsraw=None, **kwargs): +def main(args=None, argsraw=None, **kwargs) -> None: if args is None: parser = argparse.ArgumentParser( formatter_class=at.CustomArgHelpFormatter, @@ -525,7 +541,7 @@ def main(args=None, argsraw=None, **kwargs): dfnucabund, t_model_init_seconds = get_trajectory_timestepfile_nuc_abund( traj_root, particleid, "./Run_rprocess/tday_nz-plane" ) - dfnucabund.query("Z >= 1", inplace=True) + dfnucabund = dfnucabund.iloc[dfnucabund["Z"] >= 1] dfnucabund["radioactive"] = True t_model_init_days = t_model_init_seconds / (24 * 60 * 60) @@ -538,31 +554,29 @@ def main(args=None, argsraw=None, **kwargs): wollager_profilename, delim_whitespace=True, skiprows=1, names=["cellid", "velocity_outer", "rho"] ) dfdensities["cellid"] = dfdensities["cellid"].astype(int) - dfdensities["velocity_inner"] = np.concatenate(([0.0], dfdensities["velocity_outer"].values[:-1])) + dfdensities["velocity_inner"] = np.concatenate(([0.0], dfdensities["velocity_outer"].to_numpy()[:-1])) t_model_init_seconds_in = t_model_init_days_in * 24 * 60 * 60 - dfdensities.eval( + dfdensities = dfdensities.eval( ( "cellmass_grams = rho * 4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" "* (1e5 * @t_model_init_seconds_in) ** 3" ), - inplace=True, ) # now replace the density at the input time with the density at required time - dfdensities.eval( + dfdensities = dfdensities.eval( ( "rho = cellmass_grams / (" "4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" " * (1e5 * @t_model_init_seconds) ** 3)" ), - inplace=True, ) else: rho = 1e-11 print(f"{wollager_profilename} not found. Using rho {rho} g/cm3") - dfdensities = pd.DataFrame(dict(rho=rho, velocity_outer=6.0e4), index=[0]) + dfdensities = pd.DataFrame({"rho": rho, "velocity_outer": 6.0e4}, index=[0]) # print(dfdensities) @@ -571,7 +585,7 @@ def main(args=None, argsraw=None, **kwargs): dfelabundances = pd.DataFrame([dict(inputcellid=mgi + 1, **dictelemabund) for mgi in range(len(dfdensities))]) # print(dfelabundances) - at.inputmodel.save_initialabundances(dfelabundances=dfelabundances, abundancefilename=args.outputpath) + at.inputmodel.save_initelemabundances(dfelabundances=dfelabundances, abundancefilename=args.outputpath) # write model.txt diff --git a/artistools/inputmodel/rprocess_solar.py b/artistools/inputmodel/rprocess_solar.py old mode 100644 new mode 100755 index 200686ccb..47818d624 --- a/artistools/inputmodel/rprocess_solar.py +++ b/artistools/inputmodel/rprocess_solar.py @@ -8,8 +8,6 @@ import artistools as at -# import os.path - def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-outputpath", "-o", default=".", help="Path for output files") @@ -44,8 +42,7 @@ def undecayed_z(row): dfmasschain = dfbetaminus.query("A == @row.A", inplace=False) if not dfmasschain.empty: return int(dfmasschain.Z.min()) # decay to top of chain - else: - return int(row.Z) + return int(row.Z) dfsolarabund_undecayed = dfsolarabund.copy() dfsolarabund_undecayed["Z"] = dfsolarabund_undecayed.apply(undecayed_z, axis=1) @@ -62,11 +59,11 @@ def undecayed_z(row): ) normfactor = dfsolarabund_undecayed.numberfrac.sum() # convert number fractions in solar to fractions of r-process - dfsolarabund_undecayed.eval("numberfrac = numberfrac / @normfactor", inplace=True) + dfsolarabund_undecayed = dfsolarabund_undecayed.eval("numberfrac = numberfrac / @normfactor") - dfsolarabund_undecayed.eval("massfrac = numberfrac * A", inplace=True) + dfsolarabund_undecayed = dfsolarabund_undecayed.eval("massfrac = numberfrac * A") massfracnormfactor = dfsolarabund_undecayed.massfrac.sum() - dfsolarabund_undecayed.eval("massfrac = massfrac / @massfracnormfactor", inplace=True) + dfsolarabund_undecayed = dfsolarabund_undecayed.eval("massfrac = massfrac / @massfracnormfactor") # print(dfsolarabund_undecayed) @@ -80,29 +77,27 @@ def undecayed_z(row): wollager_profilename, delim_whitespace=True, skiprows=1, names=["cellid", "velocity_outer", "rho"] ) dfdensities["cellid"] = dfdensities["cellid"].astype(int) - dfdensities["velocity_inner"] = np.concatenate(([0.0], dfdensities["velocity_outer"].values[:-1])) + dfdensities["velocity_inner"] = np.concatenate(([0.0], dfdensities["velocity_outer"].to_numpy()[:-1])) t_model_init_seconds_in = t_model_init_days_in * 24 * 60 * 60 - dfdensities.eval( + dfdensities = dfdensities.eval( ( "cellmass_grams = rho * 4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" "* (1e5 * @t_model_init_seconds_in) ** 3" ), - inplace=True, ) # now replace the density at the input time with the density at required time - dfdensities.eval( + dfdensities = dfdensities.eval( ( "rho = cellmass_grams / (" "4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" " * (1e5 * @t_model_init_seconds) ** 3)" ), - inplace=True, ) else: - dfdensities = pd.DataFrame(dict(rho=10**-3, velocity_outer=6.0e4), index=[0]) + dfdensities = pd.DataFrame({"rho": 10**-3, "velocity_outer": 6.0e4}, index=[0]) # print(dfdensities) cellcount = len(dfdensities) @@ -116,7 +111,7 @@ def undecayed_z(row): dfelabundances = pd.DataFrame([dict(inputcellid=mgi + 1, **dictelemabund) for mgi in range(cellcount)]) # print(dfelabundances) - at.inputmodel.save_initialabundances(dfelabundances=dfelabundances, abundancefilename=args.outputpath) + at.inputmodel.save_initelemabundances(dfelabundances=dfelabundances, abundancefilename=args.outputpath) # write model.txt diff --git a/artistools/inputmodel/scalevelocity.py b/artistools/inputmodel/scalevelocity.py index f608288f6..7307e24b8 100755 --- a/artistools/inputmodel/scalevelocity.py +++ b/artistools/inputmodel/scalevelocity.py @@ -20,12 +20,11 @@ def addargs(parser: argparse.ArgumentParser) -> None: def eval_mshell(dfmodel: pd.DataFrame, t_model_init_seconds: float) -> None: - dfmodel.eval( + dfmodel = dfmodel.eval( ( "cellmass_grams = 10 ** logrho * 4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" "* (1e5 * @t_model_init_seconds) ** 3" ), - inplace=True, ) @@ -63,13 +62,12 @@ def main(args=None, argsraw=None, **kwargs) -> None: dfmodel.velocity_inner *= velscale dfmodel.velocity_outer *= velscale - dfmodel.eval( + dfmodel = dfmodel.eval( ( "logrho = log10(cellmass_grams / (" "4. / 3. * @math.pi * (velocity_outer ** 3 - velocity_inner ** 3)" " * (1e5 * @t_model_init_seconds) ** 3))" ), - inplace=True, ) eval_mshell(dfmodel, t_model_init_seconds) diff --git a/artistools/inputmodel/shen2018.py b/artistools/inputmodel/shen2018.py index 612b509c6..70c5f176f 100755 --- a/artistools/inputmodel/shen2018.py +++ b/artistools/inputmodel/shen2018.py @@ -8,8 +8,6 @@ import artistools as at -# import numpy as np - def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-inputpath", "-i", default="1.00_5050.dat", help="Path of input file") @@ -34,7 +32,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: for species in columns[5:]: atomic_number = at.get_atomic_number(species.rstrip("0123456789")) atomicnumberofspecies[species] = atomic_number - isotopesofelem.setdefault(atomic_number, list()).append(species) + isotopesofelem.setdefault(atomic_number, []).append(species) datain = pd.read_csv(args.inputpath, delim_whitespace=True, skiprows=0, header=[0]).dropna() @@ -94,7 +92,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: print(f'M_Ni56 = {tot_ni56mass / u.solMass.to("g"):.3f} solMass') at.save_modeldata(dfmodel, t_model_init_days, os.path.join(args.outputpath, "model.txt")) - at.inputmodel.save_initialabundances(dfelabundances, os.path.join(args.outputpath, "abundances.txt")) + at.inputmodel.save_initelemabundances(dfelabundances, os.path.join(args.outputpath, "abundances.txt")) if __name__ == "__main__": diff --git a/artistools/inputmodel/slice1Dfromconein3dmodel.py b/artistools/inputmodel/slice1Dfromconein3dmodel.py index ab77fe9fe..331df5438 100755 --- a/artistools/inputmodel/slice1Dfromconein3dmodel.py +++ b/artistools/inputmodel/slice1Dfromconein3dmodel.py @@ -9,11 +9,6 @@ import artistools as at -# import os -# import pandas as pd -# import matplotlib -# from mpl_toolkits.mplot3d import Axes3D - def make_cone(args): print("Making cone") @@ -58,9 +53,7 @@ def get_profile_along_axis(args=None, modeldata=None, derived_cols=False): # 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, dimensions=3, get_elemabundances=True, derived_cols=derived_cols - ) + modeldata, _ = at.inputmodel.get_modeldata(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" @@ -111,7 +104,7 @@ def make_1D_profile(args): if not args.positive_axis: # Invert rows and *velocity by -1 to make velocities positive for slice on negative axis - slice1D.iloc[:] = slice1D.iloc[::-1].values + slice1D.iloc[:] = slice1D.iloc[::-1].to_numpy() slice1D["vout_kmps"] = slice1D["vout_kmps"].apply(lambda x: x * -1) # with pd.option_context('display.max_rows', None, 'display.max_columns', None): diff --git a/artistools/inputmodel/test_inputmodel.py b/artistools/inputmodel/test_inputmodel.py index acf8bab6c..77c0ba81a 100644 --- a/artistools/inputmodel/test_inputmodel.py +++ b/artistools/inputmodel/test_inputmodel.py @@ -1,21 +1,9 @@ -#!/usr/bin/env python3 -# import hashlib -# import math -# import numpy as np -# import os -# import os.path -from pathlib import Path - import pandas as pd import artistools as at -import artistools.inputmodel - -# import pytest modelpath = at.get_config()["path_testartismodel"] outputpath = at.get_config()["path_testoutput"] -at.set_config("enable_diskcache", False) def test_describeinputmodel(): diff --git a/artistools/lightcurve/__init__.py b/artistools/lightcurve/__init__.py index d55341ef3..552ffcbaa 100644 --- a/artistools/lightcurve/__init__.py +++ b/artistools/lightcurve/__init__.py @@ -1,38 +1,34 @@ -#!/usr/bin/env python3 """Artistools - light curve functions.""" -from artistools.lightcurve.lightcurve import average_lightcurve_phi_bins -from artistools.lightcurve.lightcurve import average_lightcurve_theta_bins -from artistools.lightcurve.lightcurve import bolometric_magnitude -from artistools.lightcurve.lightcurve import evaluate_magnitudes -from artistools.lightcurve.lightcurve import generate_band_lightcurve_data -from artistools.lightcurve.lightcurve import get_band_lightcurve -from artistools.lightcurve.lightcurve import get_colour_delta_mag -from artistools.lightcurve.lightcurve import get_filter_data -from artistools.lightcurve.lightcurve import get_from_packets -from artistools.lightcurve.lightcurve import get_phillips_relation_data -from artistools.lightcurve.lightcurve import get_sn_sample_bol -from artistools.lightcurve.lightcurve import get_spectrum_in_filter_range -from artistools.lightcurve.lightcurve import plot_phillips_relation_data -from artistools.lightcurve.lightcurve import read_3d_gammalightcurve -from artistools.lightcurve.lightcurve import read_bol_reflightcurve_data -from artistools.lightcurve.lightcurve import read_hesma_lightcurve -from artistools.lightcurve.lightcurve import read_reflightcurve_band_data -from artistools.lightcurve.lightcurve import readfile -from artistools.lightcurve.plotlightcurve import addargs -from artistools.lightcurve.plotlightcurve import main -from artistools.lightcurve.plotlightcurve import main as plot -from artistools.lightcurve.viewingangleanalysis import calculate_costheta_phi_for_viewing_angles -from artistools.lightcurve.viewingangleanalysis import calculate_peak_time_mag_deltam15 -from artistools.lightcurve.viewingangleanalysis import get_angle_stuff -from artistools.lightcurve.viewingangleanalysis import lightcurve_polyfit -from artistools.lightcurve.viewingangleanalysis import make_peak_colour_viewing_angle_plot -from artistools.lightcurve.viewingangleanalysis import make_plot_test_viewing_angle_fit -from artistools.lightcurve.viewingangleanalysis import make_viewing_angle_risetime_peakmag_delta_m15_scatter_plot -from artistools.lightcurve.viewingangleanalysis import peakmag_risetime_declinerate_init -from artistools.lightcurve.viewingangleanalysis import plot_viewanglebrightness_at_fixed_time -from artistools.lightcurve.viewingangleanalysis import save_viewing_angle_data_for_plotting -from artistools.lightcurve.viewingangleanalysis import second_band_brightness_at_peak_first_band -from artistools.lightcurve.viewingangleanalysis import set_scatterplot_plot_params -from artistools.lightcurve.viewingangleanalysis import set_scatterplot_plotkwargs -from artistools.lightcurve.viewingangleanalysis import update_plotkwargs_for_viewingangle_colorbar -from artistools.lightcurve.viewingangleanalysis import write_viewing_angle_data +from .__main__ import main +from .lightcurve import bolometric_magnitude +from .lightcurve import evaluate_magnitudes +from .lightcurve import generate_band_lightcurve_data +from .lightcurve import get_band_lightcurve +from .lightcurve import get_colour_delta_mag +from .lightcurve import get_filter_data +from .lightcurve import get_from_packets +from .lightcurve import get_phillips_relation_data +from .lightcurve import get_sn_sample_bol +from .lightcurve import get_spectrum_in_filter_range +from .lightcurve import plot_phillips_relation_data +from .lightcurve import read_3d_gammalightcurve +from .lightcurve import read_bol_reflightcurve_data +from .lightcurve import read_hesma_lightcurve +from .lightcurve import read_reflightcurve_band_data +from .lightcurve import readfile +from .plotlightcurve import addargs +from .plotlightcurve import main as plot +from .viewingangleanalysis import calculate_peak_time_mag_deltam15 +from .viewingangleanalysis import lightcurve_polyfit +from .viewingangleanalysis import make_peak_colour_viewing_angle_plot +from .viewingangleanalysis import make_plot_test_viewing_angle_fit +from .viewingangleanalysis import make_viewing_angle_risetime_peakmag_delta_m15_scatter_plot +from .viewingangleanalysis import parse_directionbin_args +from .viewingangleanalysis import peakmag_risetime_declinerate_init +from .viewingangleanalysis import plot_viewanglebrightness_at_fixed_time +from .viewingangleanalysis import save_viewing_angle_data_for_plotting +from .viewingangleanalysis import second_band_brightness_at_peak_first_band +from .viewingangleanalysis import set_scatterplot_plot_params +from .viewingangleanalysis import set_scatterplot_plotkwargs +from .viewingangleanalysis import update_plotkwargs_for_viewingangle_colorbar +from .viewingangleanalysis import write_viewing_angle_data diff --git a/artistools/lightcurve/__main__.py b/artistools/lightcurve/__main__.py index b48f23974..a6bf0389f 100644 --- a/artistools/lightcurve/__main__.py +++ b/artistools/lightcurve/__main__.py @@ -1,5 +1,9 @@ -import artistools as at -import artistools.lightcurve.plotlightcurve +from .plotlightcurve import main as plot + + +def main() -> None: + plot() + if __name__ == "__main__": - at.lightcurve.plotlightcurve.main() + main() diff --git a/artistools/lightcurve/lightcurve.py b/artistools/lightcurve/lightcurve.py index 496613023..e388f0854 100644 --- a/artistools/lightcurve/lightcurve.py +++ b/artistools/lightcurve/lightcurve.py @@ -1,43 +1,44 @@ -#!/usr/bin/env python3 -# import glob -# import itertools import argparse import math import os from collections.abc import Collection from pathlib import Path from typing import Any +from typing import Literal from typing import Optional from typing import Union import matplotlib.pyplot as plt import numpy as np import pandas as pd +import polars as pl from astropy import constants as const from astropy import units as u import artistools as at -import artistools.spectra def readfile( filepath: Union[str, Path], - modelpath: Optional[Path] = None, - args: Union[argparse.Namespace, None] = None, -) -> Union[pd.DataFrame, dict[int, pd.DataFrame]]: - lcdata = pd.read_csv(filepath, delim_whitespace=True, header=None, names=["time", "lum", "lum_cmf"]) - - if args is not None and args.gamma and modelpath is not None and at.get_inputparams(modelpath)["n_dimensions"] == 3: - lcdata = read_3d_gammalightcurve(filepath) - - elif (args is not None and args.plotviewingangle is not None) or "res" in str(filepath): - # get a list of dfs with light curves at each viewing angle - lcdata = at.gather_res_data(lcdata, index_of_repeated_value=0) +) -> dict[int, pd.DataFrame]: + """Read an ARTIS light curve file""" + print(f"Reading {filepath}") + lcdata: dict[int, pd.DataFrame] = {} + if "_res" in str(filepath): + # get a dict of dfs with light curves at each viewing direction bin + lcdata_res = pl.read_csv( + at.zopen(filepath, "rb").read(), separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"] + ) + lcdata = at.split_dataframe_dirbins(lcdata_res, index_of_repeated_value=0, output_polarsdf=True) + else: + dfsphericalaverage = pl.read_csv( + at.zopen(filepath, "rb").read(), separator=" ", has_header=False, new_columns=["time", "lum", "lum_cmf"] + ) - elif list(lcdata.time.values) != list(sorted(lcdata.time.values)): - # the light_curve.dat file repeats x values, so keep the first half only - lcdata = lcdata.iloc[: len(lcdata) // 2] - lcdata.index.name = "timestep" + if list(dfsphericalaverage["time"].to_numpy()) != sorted(dfsphericalaverage["time"].to_numpy()): + # the light_curve.out file repeats x values, so keep the first half only + dfsphericalaverage = dfsphericalaverage[: dfsphericalaverage.height // 2] + lcdata[-1] = dfsphericalaverage return lcdata @@ -53,125 +54,136 @@ def read_3d_gammalightcurve( res_data = {} for angle in np.arange(0, 100): - res_data[angle] = lcdata[["time", angle]].copy() - res_data[angle].rename(columns={angle: "lum"}, inplace=True) + res_data[angle] = lcdata[["time", angle]] + res_data[angle] = res_data[angle].rename(columns={angle: "lum"}) return res_data def get_from_packets( - modelpath: Path, - packet_type: str = "TYPE_ESCAPE", - escape_type: str = "TYPE_RPKT", + modelpath: Union[str, Path], + escape_type: Literal["TYPE_RPKT", "TYPE_GAMMA"] = "TYPE_RPKT", maxpacketfiles: Optional[int] = None, -) -> pd.DataFrame: - import artistools.packets - - packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles=maxpacketfiles) - nprocs_read = len(packetsfiles) - assert nprocs_read > 0 + directionbins: Collection[int] = [-1], + average_over_phi: bool = False, + average_over_theta: bool = False, + get_cmf_column: bool = True, +) -> dict[int, pl.DataFrame]: + """Get ARTIS luminosity vs time from packets files""" tmidarray = at.get_timestep_times_float(modelpath=modelpath, loc="mid") timearray = at.get_timestep_times_float(modelpath=modelpath, loc="start") arr_timedelta = at.get_timestep_times_float(modelpath=modelpath, loc="delta") # timearray = np.arange(250, 350, 0.1) - _, _, vmax_cmps = at.inputmodel.get_modeldata_tuple(modelpath, getheadersonly=True, skipabundancecolumns=True) - escapesurfacegamma = math.sqrt(1 - (vmax_cmps / 29979245800) ** 2) + if get_cmf_column: + _, modelmeta = at.inputmodel.get_modeldata(modelpath, getheadersonly=True, printwarningsonly=True) + escapesurfacegamma = math.sqrt(1 - (modelmeta["vmax_cmps"] / 29979245800) ** 2) + else: + escapesurfacegamma = None timearrayplusend = np.concatenate([timearray, [timearray[-1] + arr_timedelta[-1]]]) - lcdata = pd.DataFrame( - { - "time": tmidarray, - "lum": np.zeros_like(timearray, dtype=float), - "lum_cmf": np.zeros_like(timearray, dtype=float), - } - ) - - sec_to_day = 1 / 86400 - - for packetsfile in packetsfiles: - dfpackets = at.packets.readfile(packetsfile, type=packet_type, escape_type=escape_type) - - if not (dfpackets.empty): - print(f"sum of e_cmf {dfpackets['e_cmf'].sum()} e_rf {dfpackets['e_rf'].sum()}") - - binned = pd.cut(dfpackets["t_arrive_d"], timearrayplusend, labels=False, include_lowest=True) - for binindex, e_rf_sum in dfpackets.groupby(binned)["e_rf"].sum().items(): - lcdata["lum"][binindex] += e_rf_sum - - dfpackets.eval("t_arrive_cmf_d = escape_time * @escapesurfacegamma * @sec_to_day", inplace=True) - - binned_cmf = pd.cut(dfpackets["t_arrive_cmf_d"], timearrayplusend, labels=False, include_lowest=True) - for binindex, e_cmf_sum in dfpackets.groupby(binned_cmf)["e_cmf"].sum().items(): - lcdata["lum_cmf"][binindex] += e_cmf_sum + nphibins = at.get_viewingdirection_phibincount() + ncosthetabins = at.get_viewingdirection_costhetabincount() + ndirbins = at.get_viewingdirectionbincount() - lcdata["lum"] = np.divide(lcdata["lum"] / nprocs_read * (u.erg / u.day).to("solLum"), arr_timedelta) - lcdata["lum_cmf"] = np.divide( - lcdata["lum_cmf"] / nprocs_read / escapesurfacegamma * (u.erg / u.day).to("solLum"), arr_timedelta + nprocs_read, dfpackets = at.packets.get_packets_pl( + modelpath, maxpacketfiles, packet_type="TYPE_ESCAPE", escape_type=escape_type ) - return lcdata - -def average_lightcurve_phi_bins(lcdataframes: dict[int, pd.DataFrame]) -> dict[int, pd.DataFrame]: - dirbincount = at.get_viewingdirectionbincount() - nphibins = at.get_viewingdirection_phibincount() - for start_bin in range(0, dirbincount, nphibins): - for bin_number in range(start_bin + 1, start_bin + nphibins): - lcdataframes[bin_number] = lcdataframes[bin_number].copy() # important to not affect the LRU cached copy - lcdataframes[bin_number] = lcdataframes[bin_number].set_index( - lcdataframes[start_bin].index - ) # need indexes to match or else gives NaN - lcdataframes[start_bin]["lum"] += lcdataframes[bin_number]["lum"] - del lcdataframes[bin_number] + if get_cmf_column: + dfpackets = dfpackets.with_columns( + [ + (pl.col("escape_time") * escapesurfacegamma / 86400.0).alias("t_arrive_cmf_d"), + ] + ) - lcdataframes[start_bin]["lum"] /= nphibins # every nth bin is the average of n bins - print(f"bin number {start_bin} = the average of bins {start_bin} to {start_bin + nphibins-1}") + getcols = ["t_arrive_d", "e_rf"] + if get_cmf_column: + getcols += ["e_cmf", "t_arrive_cmf_d"] + if directionbins != [-1]: + if average_over_phi: + getcols.append("costhetabin") + elif average_over_theta: + getcols.append("phibin") + else: + getcols.append("dirbin") + dfpackets = dfpackets.select(getcols).collect(streaming=True).lazy() + + lcdata = {} + for dirbin in directionbins: + if dirbin == -1: + solidanglefactor = 1.0 + pldfpackets_dirbin = dfpackets + elif average_over_phi: + assert not average_over_theta + solidanglefactor = ncosthetabins + pldfpackets_dirbin = dfpackets.filter(pl.col("costhetabin") * 10 == dirbin) + elif average_over_theta: + solidanglefactor = nphibins + pldfpackets_dirbin = dfpackets.filter(pl.col("phibin") == dirbin) + else: + solidanglefactor = ndirbins + pldfpackets_dirbin = dfpackets.filter(pl.col("dirbin") == dirbin) + + dftimebinned = at.packets.bin_and_sum( + pldfpackets_dirbin, + bincol="t_arrive_d", + bins=list(timearrayplusend), + sumcols=["e_rf"], + ) - return lcdataframes + arr_lum = ( + dftimebinned["e_rf_sum"] / nprocs_read * solidanglefactor * (u.erg / u.day).to("solLum") + ) / arr_timedelta + lcdata[dirbin] = pl.DataFrame({"time": tmidarray, "lum": arr_lum}) -def average_lightcurve_theta_bins(lcdataframes: dict[int, pd.DataFrame]) -> dict[int, pd.DataFrame]: - dirbincount = at.get_viewingdirectionbincount() - nphibins = at.get_viewingdirection_phibincount() - ncosthetabins = at.get_viewingdirection_costhetabincount() - for start_bin in range(0, nphibins): - contribbins = range(start_bin + ncosthetabins, dirbincount, ncosthetabins) - for bin_number in contribbins: - lcdataframes[bin_number] = lcdataframes[bin_number].set_index( - lcdataframes[start_bin].index - ) # need indexes to match or else gives NaN - lcdataframes[start_bin]["lum"] += lcdataframes[bin_number]["lum"] - del lcdataframes[bin_number] + if get_cmf_column: + dftimebinned_cmf = at.packets.bin_and_sum( + pldfpackets_dirbin, + bincol="t_arrive_cmf_d", + bins=list(timearrayplusend), + sumcols=["e_cmf"], + ) - lcdataframes[start_bin]["lum"] /= ncosthetabins # every nth bin is the average of n bins - print(f"bin number {start_bin} = the average of bins {[start_bin] + list(contribbins)}") + assert escapesurfacegamma is not None + lcdata[dirbin] = lcdata[dirbin].with_columns( + ( + dftimebinned_cmf["e_cmf_sum"] + / nprocs_read + * solidanglefactor + / escapesurfacegamma + * (u.erg / u.day).to("solLum") + / arr_timedelta + ).alias("lum_cmf") + ) - return lcdataframes + return lcdata def generate_band_lightcurve_data( modelpath: Path, args: argparse.Namespace, - angle: Optional[int] = None, + angle: int = -1, modelnumber: Optional[int] = None, ) -> dict: """Method adapted from https://github.com/cinserra/S3/blob/master/src/s3/SMS.py""" from scipy.interpolate import interp1d - if args and args.plotvspecpol and os.path.isfile(modelpath / "vpkt.txt"): + if args.plotvspecpol and os.path.isfile(modelpath / "vpkt.txt"): print("Found vpkt.txt, using virtual packets") - stokes_params = at.spectra.get_specpol_data(angle, modelpath) + stokes_params = ( + at.spectra.get_vspecpol_data(vspecangle=angle, modelpath=modelpath) + if angle >= 0 + else at.spectra.get_specpol_data(angle=angle, modelpath=modelpath) + ) vspecdata = stokes_params["I"] timearray = vspecdata.keys()[1:] - elif ( - args - and args.plotviewingangle - and at.anyexist(["specpol_res.out", "specpol_res.out.xz", "spec_res.out"], path=modelpath) - ): - specfilename = at.firstexisting(["specpol_res.out", "spec_res.out"], path=modelpath) + elif args.plotviewingangle and at.anyexist(["specpol_res.out", "spec_res.out"], folder=modelpath, tryzipped=True): + specfilename = at.firstexisting(["specpol_res.out", "spec_res.out"], folder=modelpath, tryzipped=True) specdataresdata = pd.read_csv(specfilename, delim_whitespace=True) - timearray = [i for i in specdataresdata.columns.values[1:] if i[-2] != "."] + timearray = [i for i in specdataresdata.columns.to_numpy()[1:] if i[-2] != "."] # elif Path(modelpath, 'specpol.out').is_file(): # specfilename = os.path.join(modelpath, "specpol.out") # specdata = pd.read_csv(specfilename, delim_whitespace=True) @@ -180,13 +192,15 @@ def generate_band_lightcurve_data( if args.plotviewingangle: print("WARNING: no direction-resolved spectra available. Using angle-averaged spectra.") - specfilename = at.firstexisting(["spec.out", "specpol.out"], path=modelpath, tryzipped=True) - if "specpol.out" in str(specfilename): - specdata = pd.read_csv(specfilename, delim_whitespace=True) - timearray = [i for i in specdata.columns.values[1:] if i[-2] != "."] # Ignore Q and U values in pol file - else: - specdata = pd.read_csv(specfilename, delim_whitespace=True) - timearray = specdata.columns.values[1:] + specfilename = at.firstexisting(["spec.out", "specpol.out"], folder=modelpath, tryzipped=True) + specdata = pd.read_csv(specfilename, delim_whitespace=True) + + timearray = ( + # Ignore Q and U values in pol file + [i for i in specdata.columns.to_numpy()[1:] if i[-2] != "."] + if "specpol.out" in str(specfilename) + else specdata.columns.to_numpy()[1:] + ) filters_dict = {} if not args.filter: @@ -194,20 +208,15 @@ def generate_band_lightcurve_data( filters_list = args.filter - res_specdata = None - if angle is not None: - try: - res_specdata = at.spectra.read_spec_res(modelpath).copy() - if args and args.average_over_phi_angle: - at.spectra.average_phi_bins(res_specdata, angle) - - except FileNotFoundError: - pass - for filter_name in filters_list: if filter_name == "bol": times, bol_magnitudes = bolometric_magnitude( - modelpath, timearray, args, angle=angle, res_specdata=res_specdata + modelpath, + timearray, + args, + angle=angle, + average_over_phi=args.average_over_phi_angle, + average_over_theta=args.average_over_theta_angle, ) filters_dict["bol"] = [ (time, bol_magnitude) @@ -230,15 +239,15 @@ def generate_band_lightcurve_data( time = float(time) if (args.timemin is None or args.timemin <= time) and (args.timemax is None or args.timemax >= time): wavelength_from_spectrum, flux = get_spectrum_in_filter_range( - modelpath, - timestep, - time, - wavefilter_min, - wavefilter_max, - angle, - res_specdata=res_specdata, - modelnumber=modelnumber, + modelpath=modelpath, + timestep=timestep, + time=time, + wavefilter_min=wavefilter_min, + wavefilter_max=wavefilter_max, + angle=angle, args=args, + average_over_phi=args.average_over_phi_angle, + average_over_theta=args.average_over_theta_angle, ) if len(wavelength_from_spectrum) > len(wavefilter): @@ -266,28 +275,31 @@ def bolometric_magnitude( modelpath: Path, timearray: Collection[float], args: argparse.Namespace, - angle: Optional[int] = None, - res_specdata: Optional[dict[int, pd.DataFrame]] = None, + angle: int = -1, + average_over_phi: bool = False, + average_over_theta: bool = False, ) -> tuple[list[float], list[float]]: magnitudes = [] times = [] for timestep, time in enumerate(timearray): time = float(time) + if (args.timemin is None or args.timemin <= time) and (args.timemax is None or args.timemax >= time): - if angle is not None: + if angle != -1: if args.plotvspecpol: spectrum = at.spectra.get_vspecpol_spectrum(modelpath, time, angle, args) else: - if res_specdata is None: - res_specdata = at.spectra.read_spec_res(modelpath) - if args and args.average_over_phi_angle: - at.spectra.average_phi_bins(res_specdata, angle) - spectrum = at.spectra.get_res_spectrum( - modelpath, timestep, timestep, angle=angle, res_specdata=res_specdata - ) + spectrum = at.spectra.get_spectrum( + modelpath=modelpath, + directionbins=[angle], + timestepmin=timestep, + timestepmax=timestep, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + )[angle] else: - spectrum = at.spectra.get_spectrum(modelpath, timestep, timestep) + spectrum = at.spectra.get_spectrum(modelpath=modelpath, timestepmin=timestep, timestepmax=timestep)[-1] integrated_flux = np.trapz(spectrum["f_lambda"], spectrum["lambda_angstroms"]) integrated_luminosity = integrated_flux * 4 * np.pi * np.power(u.Mpc.to("cm"), 2) @@ -296,7 +308,7 @@ def bolometric_magnitude( magnitudes.append(magnitude) times.append(time) # print(const.L_sun.to('erg/s').value) - # quit() + # sys.exit(1) return times, magnitudes @@ -330,22 +342,21 @@ def get_spectrum_in_filter_range( time: float, wavefilter_min: float, wavefilter_max: float, - angle: Optional[int] = None, - res_specdata: Optional[dict[int, pd.DataFrame]] = None, - modelnumber: Optional[int] = None, + angle: int = -1, spectrum: Optional[pd.DataFrame] = None, args: Optional[argparse.Namespace] = None, + average_over_phi: bool = False, + average_over_theta: bool = False, ) -> tuple[np.ndarray, np.ndarray]: - if spectrum is None: - spectrum = at.spectra.get_spectrum_at_time( - Path(modelpath), - timestep=timestep, - time=time, - args=args, - angle=angle, - res_specdata=res_specdata, - modelnumber=modelnumber, - ) + spectrum = at.spectra.get_spectrum_at_time( + Path(modelpath), + timestep=timestep, + time=time, + args=args, + dirbin=angle, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + ) wavelength_from_spectrum, flux = [], [] for wavelength, flambda in zip(spectrum["lambda_angstroms"], spectrum["f_lambda"]): @@ -359,10 +370,7 @@ def get_spectrum_in_filter_range( def evaluate_magnitudes(flux, transmission, wavelength_from_spectrum, zeropointenergyflux: float) -> float: cf = flux * transmission flux_obs = abs(np.trapz(cf, wavelength_from_spectrum)) # using trapezoidal rule to integrate - if flux_obs == 0.0: - phot_filtobs_sn = 0.0 - else: - phot_filtobs_sn = -2.5 * np.log10(flux_obs / zeropointenergyflux) + phot_filtobs_sn = 0.0 if flux_obs == 0.0 else -2.5 * np.log10(flux_obs / zeropointenergyflux) return phot_filtobs_sn @@ -391,8 +399,8 @@ def get_colour_delta_mag(band_lightcurve_data, filter_names) -> tuple[list[float time_dict_1[float(filter_1[0])] = filter_1[1] time_dict_2[float(filter_2[0])] = filter_2[1] - for time in time_dict_1.keys(): - if time in time_dict_2.keys(): # Test if time has a magnitude for both filters + for time in time_dict_1: + if time in time_dict_2: # Test if time has a magnitude for both filters plot_times.append(time) colour_delta_mag.append(time_dict_1[time] - time_dict_2[time]) @@ -423,7 +431,7 @@ def read_hesma_lightcurve(args: argparse.Namespace) -> pd.DataFrame: def read_reflightcurve_band_data(lightcurvefilename: Union[Path, str]) -> tuple[pd.DataFrame, dict[str, Any]]: filepath = Path(at.get_config()["path_artistools_dir"], "data", "lightcurves", lightcurvefilename) - metadata = at.misc.get_file_metadata(filepath) + metadata = at.get_file_metadata(filepath) data_path = os.path.join(at.get_config()["path_artistools_dir"], f"data/lightcurves/{lightcurvefilename}") lightcurve_data = pd.read_csv(data_path, comment="#") @@ -452,15 +460,16 @@ def read_reflightcurve_band_data(lightcurvefilename: Union[Path, str]) -> tuple[ def read_bol_reflightcurve_data(lightcurvefilename): - if Path(lightcurvefilename).is_file(): - data_path = Path(lightcurvefilename) - else: - data_path = Path(at.get_config()["path_artistools_dir"], "data/lightcurves/bollightcurves", lightcurvefilename) + data_path = ( + Path(lightcurvefilename) + if Path(lightcurvefilename).is_file() + else Path(at.get_config()["path_artistools_dir"], "data/lightcurves/bollightcurves", lightcurvefilename) + ) - metadata = at.misc.get_file_metadata(data_path) + metadata = at.get_file_metadata(data_path) # check for possible header line and read table - with open(data_path, "r") as flc: + with open(data_path) as flc: filepos = flc.tell() line = flc.readline() if line.startswith("#"): @@ -478,7 +487,7 @@ def read_bol_reflightcurve_data(lightcurvefilename): } if colrenames: print(f"{data_path}: renaming columns {colrenames}") - dflightcurve.rename(columns=colrenames, inplace=True) + dflightcurve = dflightcurve.rename(columns=colrenames) return dflightcurve, metadata @@ -489,7 +498,7 @@ def get_sn_sample_bol(): print(sn_data) bol_luminosity = sn_data["Lmax"].astype(float) - bol_magnitude = 4.74 - (2.5 * np.log10((10**bol_luminosity) / const.L_sun.to("erg/s").value)) # 𝑀𝑏𝑜𝑙,𝑠𝑢𝑛 = 4.74 + bol_magnitude = 4.74 - (2.5 * np.log10((10**bol_luminosity) / const.L_sun.to("erg/s").value)) # Mbol,sun = 4.74 bol_magnitude_error_upper = bol_magnitude - ( 4.74 diff --git a/artistools/lightcurve/plotlightcurve.py b/artistools/lightcurve/plotlightcurve.py index 3c61c6582..3ba265fe9 100644 --- a/artistools/lightcurve/plotlightcurve.py +++ b/artistools/lightcurve/plotlightcurve.py @@ -1,30 +1,28 @@ -#!/usr/bin/env python3 # PYTHON_ARGCOMPLETE_OK import argparse import math import multiprocessing import os +import sys from collections.abc import Iterable +from collections.abc import Sequence from pathlib import Path from typing import Any +from typing import Literal from typing import Optional +from typing import Union import argcomplete -import matplotlib +import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np import pandas as pd +import polars as pl from astropy import constants as const from extinction import apply from extinction import ccm89 import artistools as at -import artistools.plottools -import artistools.spectra - -# import glob -# import itertools -# import sys color_list = list(plt.get_cmap("tab20")(np.linspace(0, 1.0, 20))) @@ -32,9 +30,12 @@ def plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkwargs, args) -> None: axistherm.set_xscale("log") if args.plotthermalisation: - dfmodel, t_model_init_days, vmax_cmps = at.inputmodel.get_modeldata_tuple( - modelpath, skipabundancecolumns=True, derived_cols=["vel_mid_radial"] + dfmodel, modelmeta = at.inputmodel.get_modeldata( + modelpath, skipnuclidemassfraccolumns=True, derived_cols=["vel_r_mid"], dtype_backend="pyarrow" ) + + t_model_init_days = modelmeta["t_model_init_days"] + vmax_cmps = modelmeta["vmax_cmps"] model_mass_grams = dfmodel.cellmass_grams.sum() print(f" model mass: {model_mass_grams / 1.989e33:.3f} Msun") @@ -79,14 +80,12 @@ def plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkw axis.plot( depdata["tmid_days"], gammadep_lsun * 3.826e33, - **dict( - plotkwargs, - **{ - "label": plotkwargs["label"] + r" $\dot{E}_{dep,\gamma}$", - "linestyle": "dashed", - "color": color_gamma, - }, - ), + **{ + **plotkwargs, + "label": plotkwargs["label"] + r" $\dot{E}_{dep,\gamma}$", + "linestyle": "dashed", + "color": color_gamma, + }, ) color_beta = next(axis._get_lines.prop_cycler)["color"] @@ -95,30 +94,46 @@ def plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkw axis.plot( depdata["tmid_days"], depdata["eps_elec_Lsun"] * 3.826e33, - **dict( - plotkwargs, - **{ - "label": plotkwargs["label"] + r" $\dot{E}_{rad,\beta^-}$", - "linestyle": "dotted", - "color": color_beta, - }, - ), + **{ + **plotkwargs, + "label": plotkwargs["label"] + r" $\dot{E}_{rad,\beta^-}$", + "linestyle": "dotted", + "color": color_beta, + }, ) if "elecdep_Lsun" in depdata: axis.plot( depdata["tmid_days"], depdata["elecdep_Lsun"] * 3.826e33, - **dict( - plotkwargs, - **{ - "label": plotkwargs["label"] + r" $\dot{E}_{dep,\beta^-}$", - "linestyle": "dashed", - "color": color_beta, - }, - ), + **{ + **plotkwargs, + "label": plotkwargs["label"] + r" $\dot{E}_{dep,\beta^-}$", + "linestyle": "dashed", + "color": color_beta, + }, ) + c23modelpath = Path( + "/Users/luke/Library/CloudStorage/GoogleDrive-luke@lukeshingles.com/Shared" + " drives/ARTIS/artis_runs_published/Collinsetal2023-KN/sfho_long_1-35-135Msun" + ) + + c23energyrate = at.inputmodel.energyinputfiles.get_energy_rate_fromfile(c23modelpath) + c23etot, c23energydistribution_data = at.inputmodel.energyinputfiles.get_etot_fromfile(c23modelpath) + + dE = np.diff(c23energyrate["rate"] * c23etot) + dt = np.diff(c23energyrate["times"] * 24 * 60 * 60) + + axis.plot( + c23energyrate["times"][1:], + dE / dt * 0.308, + color="grey", + linestyle="--", + zorder=20, + label=r"Collins+23 $\dot{E}_{rad,\beta^-}$", + ) + # color_alpha = next(axis._get_lines.prop_cycler)['color'] color_alpha = "C1" @@ -148,20 +163,18 @@ def plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkw axistherm.plot( depdata["tmid_days"], depdata["gammadeppathint_Lsun"] / depdata["eps_gamma_Lsun"], - **dict(plotkwargs, **{"label": modelname + r" $f_\gamma$", "linestyle": "solid", "color": color_gamma}), + **{**plotkwargs, "label": modelname + r" $f_\gamma$", "linestyle": "solid", "color": color_gamma}, ) axistherm.plot( depdata["tmid_days"], depdata["elecdep_Lsun"] / depdata["eps_elec_Lsun"], - **dict( - plotkwargs, - **{ - "label": modelname + r" $f_\beta$", - "linestyle": "solid", - "color": color_beta, - }, - ), + **{ + **plotkwargs, + "label": modelname + r" $f_\beta$", + "linestyle": "solid", + "color": color_beta, + }, ) f_alpha = depdata["alphadep_Lsun"] / depdata["eps_alpha_Lsun"] @@ -172,17 +185,19 @@ def plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkw axistherm.plot( depdata["tmid_days"], f_alpha, - **dict(plotkwargs, **{"label": modelname + r" $f_\alpha$", "linestyle": "solid", "color": color_alpha}), + **{**plotkwargs, "label": modelname + r" $f_\alpha$", "linestyle": "solid", "color": color_alpha}, ) ejecta_ke: float - if "vel_mid_radial" in dfmodel.columns: - # vel_mid_radial is in cm/s - ejecta_ke = dfmodel.eval("0.5 * (cellmass_grams / 1000.) * (vel_mid_radial / 100.) ** 2").sum() + if "vel_r_mid" in dfmodel.columns: + # vel_r_mid is in cm/s + ejecta_ke = (0.5 * (dfmodel["cellmass_grams"] / 1000.0) * (dfmodel["vel_r_mid"] / 100.0) ** 2).sum() else: # velocity_inner is in km/s - ejecta_ke = dfmodel.eval("0.5 * (cellmass_grams / 1000.) * (1000. * velocity_outer) ** 2").sum() - print(f" ejecta kinetic energy: {ejecta_ke:.2e} [K] = {ejecta_ke *1e7:.2e} [erg]") + ejecta_ke = (0.5 * (dfmodel["cellmass_grams"] / 1000.0) * (1000.0 * dfmodel["velocity_outer"]) ** 2).sum() + + print(f" ejecta kinetic energy: {ejecta_ke:.2e} [J] = {ejecta_ke *1e7:.2e} [erg]") + # velocity derived from ejecta kinetic energy to match Barnes et al. (2016) Section 2.1 ejecta_v = np.sqrt(2 * ejecta_ke / (model_mass_grams * 1e-3)) v2 = ejecta_v / (0.2 * 299792458) @@ -190,44 +205,249 @@ def plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkw m5 = model_mass_grams / (5e-3 * 1.989e33) # M / (5e-3 Msun) t_ineff_gamma = 0.5 * np.sqrt(m5) / v2 - barnes_f_gamma = [1 - math.exp(-((t / t_ineff_gamma) ** -2)) for t in depdata["tmid_days"].values] + barnes_f_gamma = [1 - math.exp(-((t / t_ineff_gamma) ** -2)) for t in depdata["tmid_days"].to_numpy()] axistherm.plot( depdata["tmid_days"], barnes_f_gamma, - **dict(plotkwargs, **{"label": r"Barnes+16 $f_\gamma$", "linestyle": "dashed", "color": color_gamma}), + **{**plotkwargs, "label": r"Barnes+16 $f_\gamma$", "linestyle": "dashed", "color": color_gamma}, ) e0_beta_mev = 0.5 t_ineff_beta = 7.4 * (e0_beta_mev / 0.5) ** -0.5 * m5**0.5 * (v2 ** (-3.0 / 2)) barnes_f_beta = [ math.log(1 + 2 * (t / t_ineff_beta) ** 2) / (2 * (t / t_ineff_beta) ** 2) - for t in depdata["tmid_days"].values + for t in depdata["tmid_days"].to_numpy() ] axistherm.plot( depdata["tmid_days"], barnes_f_beta, - **dict(plotkwargs, **{"label": r"Barnes+16 $f_\beta$", "linestyle": "dashed", "color": color_beta}), + **{**plotkwargs, "label": r"Barnes+16 $f_\beta$", "linestyle": "dashed", "color": color_beta}, ) e0_alpha_mev = 6.0 t_ineff_alpha = 4.3 * 1.8 * (e0_alpha_mev / 6.0) ** -0.5 * m5**0.5 * (v2 ** (-3.0 / 2)) barnes_f_alpha = [ math.log(1 + 2 * (t / t_ineff_alpha) ** 2) / (2 * (t / t_ineff_alpha) ** 2) - for t in depdata["tmid_days"].values + for t in depdata["tmid_days"].to_numpy() ] axistherm.plot( depdata["tmid_days"], barnes_f_alpha, - **dict(plotkwargs, **{"label": r"Barnes+16 $f_\alpha$", "linestyle": "dashed", "color": color_alpha}), + **{**plotkwargs, "label": r"Barnes+16 $f_\alpha$", "linestyle": "dashed", "color": color_alpha}, + ) + + +def plot_artis_lightcurve( + modelpath: Union[str, Path], + axis, + lcindex: int = 0, + label: Optional[str] = None, + escape_type: Literal["TYPE_RPKT", "TYPE_GAMMA"] = "TYPE_RPKT", + frompackets: bool = False, + maxpacketfiles: Optional[int] = None, + axistherm=None, + directionbins: Sequence[int] = [-1], + average_over_phi: bool = False, + average_over_theta: bool = False, + args=None, +) -> Optional[pd.DataFrame]: + lcfilename = None + modelpath = Path(modelpath) + if Path(modelpath).is_file(): # handle e.g. modelpath = 'modelpath/light_curve.out' + lcfilename = Path(modelpath).parts[-1] + modelpath = Path(modelpath).parent + + if not modelpath.is_dir(): + print(f"WARNING: Skipping because {modelpath} does not exist") + return None + + modelname = at.get_model_name(modelpath) if label is None else label + if lcfilename is not None: + modelname += f" {lcfilename}" + if not modelname: + print("====> (no series label)") + else: + print(f"====> {modelname}") + print(f" folder: {modelpath.resolve().parts[-1]}") + + if args is not None and args.title: + axis.set_title(modelname) + + if directionbins is None: + directionbins = [-1] + + if frompackets: + lcdataframes = at.lightcurve.get_from_packets( + modelpath, + escape_type=escape_type, + maxpacketfiles=maxpacketfiles, + directionbins=directionbins, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + get_cmf_column=args.plotcmf, + ) + else: + if lcfilename is None: + lcfilename = ( + "light_curve_res.out" + if directionbins != [-1] + else "gamma_light_curve.out" if escape_type == "TYPE_GAMMA" else "light_curve.out" + ) + + try: + lcpath = at.firstexisting(lcfilename, folder=modelpath, tryzipped=True) + except FileNotFoundError: + print(f"WARNING: Skipping {modelname} because {lcfilename} does not exist") + return None + + lcdataframes = at.lightcurve.readfile(lcpath) + + if average_over_phi: + lcdataframes = at.average_direction_bins(lcdataframes, overangle="phi") + + if average_over_theta: + lcdataframes = at.average_direction_bins(lcdataframes, overangle="theta") + + plotkwargs: dict[str, Any] = {} + plotkwargs["label"] = modelname + plotkwargs["linestyle"] = args.linestyle[lcindex] + plotkwargs["color"] = args.color[lcindex] + if args.dashes[lcindex]: + plotkwargs["dashes"] = args.dashes[lcindex] + if args.linewidth[lcindex]: + plotkwargs["linewidth"] = args.linewidth[lcindex] + + # check if doing viewing angle stuff, and if so define which data to use + dirbins, angle_definition = at.lightcurve.parse_directionbin_args(modelpath, args) + + if args.colorbarcostheta or args.colorbarphi: + costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_costhetabin_phibin_labels() + scaledmap = make_colorbar_viewingangles_colormap() + + lctimemin, lctimemax = float(lcdataframes[dirbins[0]]["time"].to_numpy().min()), float( + lcdataframes[dirbins[0]]["time"].to_numpy().max() + ) + + print(f" range of light curve: {lctimemin:.2f} to {lctimemax:.2f} days") + try: + nts_last, validrange_start_days, validrange_end_days = at.get_escaped_arrivalrange(modelpath) + if validrange_start_days is not None and validrange_end_days is not None: + str_valid_range = f"{validrange_start_days:.2f} to {validrange_end_days:.2f} days" + else: + str_valid_range = f"{validrange_start_days} to {validrange_end_days} days" + print(f" range of validity (last timestep {nts_last}): {str_valid_range}") + except FileNotFoundError: + print( + " range of validity: could not determine due to missing files " + "(requires deposition.out, input.txt, model.txt)" ) + nts_last, validrange_start_days, validrange_end_days = None, float("-inf"), float("inf") + + for dirbin in dirbins: + lcdata = lcdataframes[dirbin] + + if dirbin != -1: + print(f" directionbin {dirbin:4d} {angle_definition[dirbin]}") + + if args.colorbarcostheta or args.colorbarphi: + plotkwargs["alpha"] = 0.75 + plotkwargs["label"] = None + # Update plotkwargs with viewing angle colour + plotkwargs, colorindex = get_viewinganglecolor_for_colorbar( + dirbin, + costheta_viewing_angle_bins, + phi_viewing_angle_bins, + scaledmap, + plotkwargs, + args, + ) + if args.average_over_phi_angle: + plotkwargs["color"] = "lightgrey" + else: + # the first dirbin should use the color argument (which has been removed from the color cycle) + if dirbin != dirbins[0]: + plotkwargs["color"] = None + plotkwargs["label"] = ( + f"{modelname}\n{angle_definition[dirbin]}" if modelname else angle_definition[dirbin] + ) + + filterfunc = at.get_filterfunc(args) + if filterfunc is not None: + lcdata = lcdata.with_columns( + pl.from_numpy(filterfunc(lcdata["lum"].to_numpy()), schema=["lum"]).get_column("lum") + ) + + if not args.Lsun or args.magnitude: + # convert luminosity from Lsun to erg/s + lcdata = lcdata.with_columns(pl.col("lum") * 3.826e33) + if "lum_cmf" in lcdata.columns: + lcdata = lcdata.with_columns(pl.col("lum_cmf") * 3.826e33) + + if args.magnitude: + # convert to bol magnitude + lcdata["mag"] = 4.74 - (2.5 * np.log10(lcdata["lum"] / const.L_sun.to("erg/s").value)) + ycolumn = "mag" + else: + ycolumn = "lum" + + if ( + args.average_over_phi_angle + and dirbin % at.get_viewingdirection_costhetabincount() == 0 + and (args.colorbarcostheta or args.colorbarphi) + ): + plotkwargs["color"] = scaledmap.to_rgba(colorindex) # Update colours for light curves averaged over phi + plotkwargs["zorder"] = 10 + + # show the parts of the light curve that are outside the valid arrival range as partially transparent + if validrange_start_days is None or validrange_end_days is None: + # entire range is invalid + lcdata_before_valid = lcdata + lcdata_after_valid = pd.DataFrame(data=None, columns=lcdata.columns) + lcdata_valid = pd.DataFrame(data=None, columns=lcdata.columns) + else: + lcdata_valid = lcdata.filter( + (pl.col("time") >= validrange_start_days) & (pl.col("time") <= validrange_end_days) + ) + + lcdata_before_valid = lcdata.filter(pl.col("time") >= lcdata_valid["time"].min()) + lcdata_after_valid = lcdata.filter(pl.col("time") >= lcdata_valid["time"].max()) + + axis.plot(lcdata_valid["time"], lcdata_valid[ycolumn], **plotkwargs) + + if args.plotinvalidpart: + plotkwargs_invalidrange = plotkwargs.copy() + plotkwargs_invalidrange.update({"label": None, "alpha": 0.5}) + axis.plot(lcdata_before_valid["time"], lcdata_before_valid[ycolumn], **plotkwargs_invalidrange) + axis.plot(lcdata_after_valid["time"], lcdata_after_valid[ycolumn], **plotkwargs_invalidrange) + + if args.print_data: + print(lcdata[["time", ycolumn, "lum_cmf"]]) + if args.plotcmf: + plotkwargs["linewidth"] = 1 + plotkwargs["label"] += " (cmf)" + plotkwargs["linestyle"] = "dashed" + # plotkwargs['color'] = 'tab:orange' + axis.plot(lcdata["time"], lcdata["lum_cmf"], **plotkwargs) -def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type=False, maxpacketfiles=None, args=None): - """Use light_curve.out or light_curve_res.out files to plot light curve""" + if args.plotdeposition or args.plotthermalisation: + plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkwargs, args) + return lcdataframes + + +def make_lightcurve_plot( + modelpaths: Sequence[Union[str, Path]], + filenameout: str, + frompackets: bool = False, + escape_type: Literal["TYPE_RPKT", "TYPE_GAMMA"] = "TYPE_RPKT", + maxpacketfiles: Optional[int] = None, + args=None, +): + """Use light_curve.out or light_curve_res.out files to plot light curve.""" fig, axis = plt.subplots( nrows=1, ncols=1, @@ -247,13 +467,6 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type else: axistherm = None - if not frompackets and escape_type not in ["TYPE_RPKT", "TYPE_GAMMA"]: - print(f"Escape_type of {escape_type} not one of TYPE_RPKT or TYPE_GAMMA, so frompackets must be enabled") - assert False - elif not frompackets and args.packet_type != "TYPE_ESCAPE" and args.packet_type is not None: - print("Looking for non-escaped packets, so frompackets must be enabled") - assert False - # take any assigned colours our of the cycle colors = [ color for i, color in enumerate(plt.rcParams["axes.prop_cycle"].by_key()["color"]) if f"C{i}" not in args.color @@ -261,14 +474,15 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type axis.set_prop_cycle(color=colors) reflightcurveindex = 0 - for seriesindex, modelpath in enumerate(modelpaths): + plottedsomething = False + for lcindex, modelpath in enumerate(modelpaths): if not Path(modelpath).is_dir() and not Path(modelpath).exists() and "." in str(modelpath): bolreflightcurve = Path(modelpath) dflightcurve, metadata = at.lightcurve.read_bol_reflightcurve_data(bolreflightcurve) lightcurvelabel = metadata.get("label", bolreflightcurve) color = ["0.0", "0.5", "0.7"][reflightcurveindex] - plotkwargs = dict(label=lightcurvelabel, color=color, zorder=0) + plotkwargs = {"label": lightcurvelabel, "color": color, "zorder": 0} if ( "luminosity_errminus_erg/s" in dflightcurve.columns and "luminosity_errplus_erg/s" in dflightcurve.columns @@ -285,156 +499,29 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type axis.scatter(dflightcurve["time_days"], dflightcurve["luminosity_erg/s"], **plotkwargs) print(f"====> {lightcurvelabel}") reflightcurveindex += 1 - continue - - lcname = None - if not Path(modelpath).is_dir(): # handle e.g. modelpath/light_curve.out - lcname = Path(modelpath).parts[-1] - modelpath = Path(modelpath).parent - modelname = at.get_model_name(modelpath) if args.label[seriesindex] is None else args.label[seriesindex] - if lcname is not None: - modelname += f" {lcname}" - print(f"====> {modelname}") - - if lcname is None: - lcname = "gamma_light_curve.out" if (escape_type == "TYPE_GAMMA" and not frompackets) else "light_curve.out" - if args.plotviewingangle is not None and lcname == "light_curve.out": - lcname = "light_curve_res.out" - try: - lcpath = at.firstexisting(lcname, path=modelpath, tryzipped=True) - except FileNotFoundError: - print(f"Skipping {modelname} because {lcname} does not exist") - continue - if not os.path.exists(str(lcpath)): - print(f"Skipping {modelname} because {lcpath} does not exist") - continue - elif frompackets: - lcdata = at.lightcurve.get_from_packets( - modelpath, packet_type=args.packet_type, escape_type=escape_type, maxpacketfiles=maxpacketfiles - ) + plottedsomething = True else: - lcdata = at.lightcurve.readfile(lcpath, modelpath, args) - - plotkwargs = {} - plotkwargs["label"] = modelname - plotkwargs["linestyle"] = args.linestyle[seriesindex] - plotkwargs["color"] = args.color[seriesindex] - if args.dashes[seriesindex]: - plotkwargs["dashes"] = args.dashes[seriesindex] - if args.linewidth[seriesindex]: - plotkwargs["linewidth"] = args.linewidth[seriesindex] - - if args.plotdeposition or args.plotthermalisation: - plot_deposition_thermalisation(axis, axistherm, modelpath, modelname, plotkwargs, args) - - # check if doing viewing angle stuff, and if so define which data to use - angles, viewing_angles, angle_definition = at.lightcurve.get_angle_stuff(modelpath, args) - if args.plotviewingangle: - lcdataframes = lcdata - if args.average_over_phi_angle: - lcdataframes = at.lightcurve.average_lightcurve_phi_bins(lcdataframes) - - if args.average_over_theta_angle: - lcdataframes = at.lightcurve.average_lightcurve_theta_bins(lcdataframes) - - if args.colorbarcostheta or args.colorbarphi: - costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_viewinganglebin_definitions() - scaledmap = make_colorbar_viewingangles_colormap() - - if args.plotviewingangle: - lcdata = lcdataframes[0] - print(f" range of light curve: {lcdata.time.min():.2f} to {lcdata.time.max():.2f} days") - try: - nts_last, validrange_start_days, validrange_end_days = at.get_escaped_arrivalrange(modelpath) - if validrange_start_days is not None and validrange_end_days is not None: - str_valid_range = f"{validrange_start_days:.2f} to {validrange_end_days:.2f} days" - else: - str_valid_range = f"{validrange_start_days} to {validrange_end_days} days" - print(f" range of validity (last timestep {nts_last}): {str_valid_range}") - except FileNotFoundError: - print( - " range of validity: could not determine due to missing files " - "(requires deposition.out, input.txt, model.txt)" + lcdataframes = plot_artis_lightcurve( + modelpath=modelpath, + lcindex=lcindex, + label=args.label[lcindex], + axis=axis, + escape_type=escape_type, + frompackets=frompackets, + maxpacketfiles=maxpacketfiles, + axistherm=axistherm, + directionbins=args.plotviewingangle if args.plotviewingangle is not None else [-1], + average_over_phi=args.average_over_phi_angle, + average_over_theta=args.average_over_theta_angle, + args=args, ) - nts_last, validrange_start_days, validrange_end_days = None, float("-inf"), float("inf") - - for angleindex, angle in enumerate(angles): - if args.plotviewingangle: - lcdata = lcdataframes[angle] - - if args.colorbarcostheta or args.colorbarphi: - plotkwargs["alpha"] = 0.75 - plotkwargs["label"] = None - # Update plotkwargs with viewing angle colour - plotkwargs, colorindex = get_viewinganglecolor_for_colorbar( - angle_definition, - angle, - costheta_viewing_angle_bins, - phi_viewing_angle_bins, - scaledmap, - plotkwargs, - args, - ) - if args.average_over_phi_angle: # if angles plotted that are not averaged over phi - plotkwargs["color"] = "lightgrey" # then plot these in grey - else: - plotkwargs["color"] = None - plotkwargs["label"] = ( - f"{modelname}\n{angle_definition[angle]}" if modelname else angle_definition[angle] - ) - - filterfunc = at.get_filterfunc(args) - if filterfunc is not None: - lcdata["lum"] = filterfunc(lcdata["lum"]) - - if not args.Lsun or args.magnitude: - # convert luminosity from Lsun to erg/s - lcdata.eval("lum = lum * 3.826e33", inplace=True) - lcdata.eval("lum_cmf = lum_cmf * 3.826e33", inplace=True) - - if args.magnitude: - # convert to bol magnitude - lcdata["mag"] = 4.74 - (2.5 * np.log10(lcdata["lum"] / const.L_sun.to("erg/s").value)) - axis.plot(lcdata["time"], lcdata["mag"], **plotkwargs) - else: - # show the parts of the light curve that are outside the valid arrival range partially transparent - if validrange_start_days is None or validrange_end_days is None: - # entire range is invalid - lcdata_before_valid = lcdata - lcdata_after_valid = pd.DataFrame(data=None, columns=lcdata.columns) - lcdata_valid = pd.DataFrame(data=None, columns=lcdata.columns) - else: - lcdata_valid = lcdata.query( - "time >= @validrange_start_days and time <= @validrange_end_days", inplace=False - ) - lcdata_before_valid = lcdata.query("time <= @lcdata_valid.time.min()") - lcdata_after_valid = lcdata.query("time >= @lcdata_valid.time.max()") - - if args.average_over_phi_angle and angle % 10 == 0 and (args.colorbarcostheta or args.colorbarphi): - color = scaledmap.to_rgba(colorindex) # Update colours for light curves averaged over phi - axis.plot(lcdata["time"], lcdata["lum"], color=color, zorder=10) - else: - axis.plot(lcdata_valid["time"], lcdata_valid["lum"], **plotkwargs) - if args.plotinvalidpart: - plotkwargs_invalidrange = plotkwargs.copy() - plotkwargs_invalidrange.update({"label": None, "alpha": 0.5}) - axis.plot(lcdata_before_valid["time"], lcdata_before_valid["lum"], **plotkwargs_invalidrange) - axis.plot(lcdata_after_valid["time"], lcdata_after_valid["lum"], **plotkwargs_invalidrange) - - if args.print_data: - print(lcdata[["time", "lum", "lum_cmf"]].to_string(index=False)) - if args.plotcmf: - plotkwargs["linewidth"] = 1 - plotkwargs["label"] += " (cmf)" - plotkwargs["linestyle"] = "dashed" - # plotkwargs['color'] = 'tab:orange' - axis.plot(lcdata.time, lcdata["lum_cmf"], **plotkwargs) + plottedsomething = plottedsomething or (lcdataframes is not None) if args.reflightcurves: for bolreflightcurve in args.reflightcurves: if args.Lsun: print("Check units - trying to plot ref light curve in erg/s") - quit() + sys.exit(1) bollightcurve_data, metadata = at.lightcurve.read_bol_reflightcurve_data(bolreflightcurve) axis.scatter( bollightcurve_data["time_days"], @@ -442,9 +529,12 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type label=metadata.get("label", bolreflightcurve), color="k", ) + plottedsomething = True + + assert plottedsomething if args.magnitude: - plt.gca().invert_yaxis() + axis.invert_yaxis() if args.xmin is not None: axis.set_xlim(left=args.xmin) @@ -466,10 +556,7 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type if args.magnitude: axis.set_ylabel("Absolute Bolometric Magnitude") else: - if not args.Lsun: - str_units = " [erg/s]" - else: - str_units = r"$/ \mathrm{L}_\odot$" + str_units = " [erg/s]" if not args.Lsun else "$/ \\mathrm{L}_\\odot$" if args.plotdeposition: yvarname = "" elif escape_type == "TYPE_GAMMA": @@ -480,10 +567,9 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type yvarname = r"$\mathrm{L}_{\mathrm{" + escape_type.replace("_", r"\_") + r"}}$" axis.set_ylabel(yvarname + str_units) - if args.title: - axis.set_title(modelname) - if args.colorbarcostheta or args.colorbarphi: + costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_costhetabin_phibin_labels() + scaledmap = make_colorbar_viewingangles_colormap() make_colorbar_viewingangles(phi_viewing_angle_bins, scaledmap, args) if args.logscalex: @@ -521,7 +607,7 @@ def make_lightcurve_plot(modelpaths, filenameout, frompackets=False, escape_type def create_axes(args): if "labelfontsize" in args: font = {"size": args.labelfontsize} - matplotlib.rc("font", **font) + mpl.rc("font", **font) args.subplots = False # todo: set as command line arg @@ -565,16 +651,9 @@ def get_linelabel( angle_definition: Optional[dict[int, str]], args, ) -> str: - if args.plotvspecpol and angle is not None and os.path.isfile(modelpath / "vpkt.txt"): - vpkt_config = at.get_vpkt_config(modelpath) - viewing_angle = round(math.degrees(math.acos(vpkt_config["cos_theta"][angle]))) - linelabel = rf"$\theta$ = {viewing_angle}" # todo: update to be consistent with res definition - elif args.plotviewingangle and angle is not None and os.path.isfile(modelpath / "specpol_res.out"): + if angle is not None and angle != -1: assert angle_definition is not None - if args.nomodelname: - linelabel = f"{angle_definition[angle]}" - else: - linelabel = f"{modelname} {angle_definition[angle]}" + linelabel = f"{angle_definition[angle]}" if args.nomodelname else f"{modelname} {angle_definition[angle]}" # linelabel = None # linelabel = fr"{modelname} $\theta$ = {angle_names[index]}$^\circ$" # plt.plot(time, magnitude, label=linelabel, linewidth=3) @@ -626,30 +705,35 @@ def set_lightcurve_plot_labels(fig, ax, filternames_conversion_dict, args, band_ ax.set_xlabel("Time Since Explosion [days]", fontsize=args.labelfontsize) if ylabel is None: print("failed to set ylabel") - quit() + sys.exit(1) return fig, ax def make_colorbar_viewingangles_colormap(): - norm = matplotlib.colors.Normalize(vmin=0, vmax=9) - scaledmap = matplotlib.cm.ScalarMappable(cmap="tab10", norm=norm) + norm = mpl.colors.Normalize(vmin=0, vmax=9) + scaledmap = mpl.cm.ScalarMappable(cmap="tab10", norm=norm) scaledmap.set_array([]) return scaledmap def get_viewinganglecolor_for_colorbar( - angle_definition, angle, costheta_viewing_angle_bins, phi_viewing_angle_bins, scaledmap, plotkwargs, args + angle: int, costheta_viewing_angle_bins, phi_viewing_angle_bins, scaledmap, plotkwargs, args ): + nphibins = at.get_viewingdirection_phibincount() if args.colorbarcostheta: - colorindex = costheta_viewing_angle_bins.index(angle_definition[angle].split(", ")[0]) + costheta_index = angle // nphibins + colorindex = costheta_index plotkwargs["color"] = scaledmap.to_rgba(colorindex) if args.colorbarphi: - colorindex = phi_viewing_angle_bins.index(angle_definition[angle].split(", ")[1]) + phi_index = angle % nphibins + colorindex = phi_index + assert nphibins == 10 reorderphibins = {5: 9, 6: 8, 7: 7, 8: 6, 9: 5} print("Reordering phi bins") - if colorindex in reorderphibins.keys(): + if colorindex in reorderphibins: colorindex = reorderphibins[colorindex] plotkwargs["color"] = scaledmap.to_rgba(colorindex) + return plotkwargs, colorindex @@ -658,31 +742,31 @@ def make_colorbar_viewingangles(phi_viewing_angle_bins, scaledmap, args, fig=Non # ticklabels = costheta_viewing_angle_bins ticklabels = [" -1", " -0.8", " -0.6", " -0.4", " -0.2", " 0", " 0.2", " 0.4", " 0.6", " 0.8", " 1"] ticklocs = np.linspace(0, 9, num=11) - label = "cos(\u03b8)" + label = "cos θ" if args.colorbarphi: print("reordered phi bins") phi_viewing_angle_bins_reordered = [ "0", - "\u03c0/5", - "2\u03c0/5", - "3\u03c0/5", - "4\u03c0/5", - "\u03c0", - "6\u03c0/5", - "7\u03c0/5", - "8\u03c0/5", - "9\u03c0/5", - "2\u03c0", + "π/5", + "2π/5", + "3π/5", + "4π/5", + "π", + "6π/5", + "7π/5", + "8π/5", + "9π/5", + "2π", ] ticklabels = phi_viewing_angle_bins_reordered # ticklabels = phi_viewing_angle_bins ticklocs = np.linspace(0, 9, num=11) - label = "\u03d5 bin" + label = "ϕ bin" hidecolorbar = False if not hidecolorbar: if fig: - from mpl_toolkits.axes_grid1.inset_locator import inset_axes + # from mpl_toolkits.axes_grid1.inset_locator import inset_axes # cax = plt.axes([0.3, 0.97, 0.45, 0.02]) #2nd and 4th move up and down. 1st left and right. 3rd bar width cax = plt.axes([0.2, 0.98, 0.65, 0.04]) @@ -691,8 +775,8 @@ def make_colorbar_viewingangles(phi_viewing_angle_bins, scaledmap, args, fig=Non cbar = plt.colorbar(scaledmap) if label: cbar.set_label(label, rotation=0) - cbar.locator = matplotlib.ticker.FixedLocator(ticklocs) - cbar.formatter = matplotlib.ticker.FixedFormatter(ticklabels) + cbar.locator = mpl.ticker.FixedLocator(ticklocs) + cbar.formatter = mpl.ticker.FixedFormatter(ticklabels) cbar.update_ticks() @@ -706,14 +790,15 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo plotkwargs: dict[str, Any] = {} if args.colorbarcostheta or args.colorbarphi: - costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_viewinganglebin_definitions() + costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_costhetabin_phibin_labels() scaledmap = make_colorbar_viewingangles_colormap() + first_band_name = None for modelnumber, modelpath in enumerate(modelpaths): modelpath = Path(modelpath) # Make sure modelpath is defined as path. May not be necessary # check if doing viewing angle stuff, and if so define which data to use - angles, viewing_angles, angle_definition = at.lightcurve.get_angle_stuff(modelpath, args) + angles, angle_definition = at.lightcurve.parse_directionbin_args(modelpath, args) for index, angle in enumerate(angles): modelname = at.get_model_name(modelpath) @@ -727,6 +812,8 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo plotkwargs["label"] = str(args.plot_hesma_model).split("_")[:3] for plotnumber, band_name in enumerate(band_lightcurve_data): + if first_band_name is None: + first_band_name = band_name time, brightness_in_mag = at.lightcurve.get_band_lightcurve(band_lightcurve_data, band_name, args) if args.print_data or args.write_data: @@ -738,10 +825,11 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo txtlinesout.append(f"{t} {m}") txtout = "\n".join(txtlinesout) if args.write_data: - if angle is not None: - bandoutfile = Path(f"band_{band_name}_angle_{angle}.txt") - else: - bandoutfile = Path(f"band_{band_name}.txt") + bandoutfile = ( + Path(f"band_{band_name}_angle_{angle}.txt") + if angle != -1 + else Path(f"band_{band_name}.txt") + ) with bandoutfile.open("w") as f: f.write(txtout) print(f"Saved {bandoutfile}") @@ -760,16 +848,13 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo # global define_colours_list # plt.plot(time, brightness_in_mag, label=modelname, color=define_colours_list[angle], linewidth=3) - if ( - modelnumber == 0 and args.plot_hesma_model and band_name in hesma_model.keys() - ): # todo: see if this works + if modelnumber == 0 and args.plot_hesma_model and band_name in hesma_model: # todo: see if this works ax.plot(hesma_model.t, hesma_model[band_name], color="black") # axarr[plotnumber].axis([0, 60, -16, -19.5]) - if band_name in filternames_conversion_dict: - text_key = filternames_conversion_dict[band_name] - else: - text_key = band_name + text_key = ( + filternames_conversion_dict[band_name] if band_name in filternames_conversion_dict else band_name + ) if args.subplots: ax[plotnumber].annotate( @@ -793,7 +878,7 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo define_colours_list = args.refspeccolors markers = args.refspecmarkers for i, reflightcurve in enumerate(args.reflightcurves): - plot_lightcurve_from_data( + plot_lightcurve_from_refdata( band_lightcurve_data.keys(), reflightcurve, define_colours_list[i], @@ -803,16 +888,16 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo plotnumber, ) - if args.color: - plotkwargs["color"] = args.color[modelnumber] - else: - plotkwargs["color"] = define_colours_list[modelnumber] + if len(angles) == 1: + if args.color: + plotkwargs["color"] = args.color[modelnumber] + else: + plotkwargs["color"] = define_colours_list[modelnumber] if args.colorbarcostheta or args.colorbarphi: # Update plotkwargs with viewing angle colour plotkwargs["label"] = None plotkwargs, _ = get_viewinganglecolor_for_colorbar( - angle_definition, angle, costheta_viewing_angle_bins, phi_viewing_angle_bins, @@ -826,7 +911,6 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo # if not (args.test_viewing_angle_fit or args.calculate_peak_time_mag_deltam15_bool): curax = ax[plotnumber] if args.subplots else ax - curax.invert_yaxis() if args.subplots: if len(angles) > 1 or (args.plotviewingangle and os.path.isfile(modelpath / "specpol_res.out")): ax[plotnumber].plot(time, brightness_in_mag, linewidth=4, **plotkwargs) @@ -837,21 +921,26 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo # ax[plotnumber].plot( # cmfgen_mags['time[d]'], cmfgen_mags[key], label='CMFGEN', color='k', linewidth=3) else: - ax.plot(time, brightness_in_mag, linewidth=3.5, **plotkwargs) # color=color, linestyle=linestyle) + curax.plot( + time, brightness_in_mag, linewidth=3.5, **plotkwargs + ) # color=color, linestyle=linestyle) - import artistools.plottools + at.set_mpl_style() ax = at.plottools.set_axis_properties(ax, args) - fig, ax = set_lightcurve_plot_labels(fig, ax, filternames_conversion_dict, args, band_name=band_name) + fig, ax = set_lightcurve_plot_labels(fig, ax, filternames_conversion_dict, args, band_name=first_band_name) set_lightcurveplot_legend(ax, args) if args.colorbarcostheta or args.colorbarphi: make_colorbar_viewingangles(phi_viewing_angle_bins, scaledmap, args, fig=fig, ax=ax) if args.filter and len(band_lightcurve_data) == 1: - args.outputfile = os.path.join(outputfolder, f"plot{band_name}lightcurves.pdf") + args.outputfile = os.path.join(outputfolder, f"plot{first_band_name}lightcurves.pdf") if args.show: plt.show() + + (ax[0] if args.subplots else ax).invert_yaxis() + plt.savefig(args.outputfile, format="pdf") print(f"Saved figure: {args.outputfile}") @@ -881,7 +970,7 @@ def make_band_lightcurves_plot(modelpaths, filternames_conversion_dict, outputfo # colours = args.refspeccolors # markers = args.refspecmarkers # for i, reflightcurve in enumerate(args.reflightcurves): -# plot_lightcurve_from_data(filters_dict.keys(), reflightcurve, colours[i], markers[i], +# plot_lightcurve_from_refdata(filters_dict.keys(), reflightcurve, colours[i], markers[i], # filternames_conversion_dict) @@ -898,7 +987,7 @@ def colour_evolution_plot(modelpaths, filternames_conversion_dict, outputfolder, modelname = at.get_model_name(modelpath) print(f"Reading spectra: {modelname}") - angles, viewing_angles, angle_definition = at.lightcurve.get_angle_stuff(modelpath, args) + angles, angle_definition = at.lightcurve.parse_directionbin_args(modelpath, args) for index, angle in enumerate(angles): for plotnumber, filters in enumerate(args.colour_evolution): @@ -1003,7 +1092,7 @@ def colour_evolution_plot(modelpaths, filternames_conversion_dict, outputfolder, # color='k' -def plot_lightcurve_from_data( +def plot_lightcurve_from_refdata( filter_names, lightcurvefilename, color, marker, filternames_conversion_dict, ax, plotnumber ): lightcurve_data, metadata = at.lightcurve.read_reflightcurve_band_data(lightcurvefilename) @@ -1014,13 +1103,13 @@ def plot_lightcurve_from_data( for plotnumber, filter_name in enumerate(filter_names): if filter_name == "bol": continue - f = open(filterdir / Path(f"{filter_name}.txt")) + f = filterdir / Path(f"{filter_name}.txt").open() lines = f.readlines() lambda0 = float(lines[2]) if filter_name == "bol": continue - elif filter_name in filternames_conversion_dict: + if filter_name in filternames_conversion_dict: filter_name = filternames_conversion_dict[filter_name] filter_data[filter_name] = lightcurve_data.loc[lightcurve_data["band"] == filter_name] # plt.plot(limits_x, limits_y, 'v', label=None, color=color) @@ -1092,7 +1181,7 @@ def plot_color_evolution_from_data( filter_data = [] for i, filter_name in enumerate(filter_names): - f = open(filterdir / Path(f"{filter_name}.txt")) + f = filterdir / Path(f"{filter_name}.txt").open() lines = f.readlines() lambda0 = float(lines[2]) @@ -1182,8 +1271,6 @@ def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("--gamma", action="store_true", help="Make light curve from gamma rays instead of R-packets") - parser.add_argument("-packet_type", default="TYPE_ESCAPE", help="Type of escaping packets") - parser.add_argument("-escape_type", default="TYPE_RPKT", help="Type of escaping packets") parser.add_argument("-o", "-outputfile", action="store", dest="outputfile", type=Path, help="Filename for PDF file") @@ -1288,7 +1375,7 @@ def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument( "-filtersavgol", nargs=2, - help="Savitzky–Golay filter. Specify the window_length and poly_order.e.g. -filtersavgol 5 3", + help="Savitzky-Golay filter. Specify the window_length and poly_order.e.g. -filtersavgol 5 3", ) parser.add_argument( @@ -1452,11 +1539,12 @@ def main(args=None, argsraw=None, **kwargs): print("WARNING: --average_every_tenth_viewing_angle is deprecated. use --average_over_phi_angle instead") args.average_over_phi_angle = True - if not args.modelpath and not args.colour_evolution: - args.modelpath = ["."] - elif not args.modelpath and (args.filter or args.colour_evolution): + at.set_mpl_style() + + if not args.modelpath: args.modelpath = ["."] - elif not isinstance(args.modelpath, Iterable): + + if not isinstance(args.modelpath, Iterable): args.modelpath = [args.modelpath] args.modelpath = at.flatten_list(args.modelpath) @@ -1489,9 +1577,9 @@ def main(args=None, argsraw=None, **kwargs): if not args.outputfile: outputfolder = Path() args.outputfile = defaultoutputfile - elif os.path.isdir(args.outputfile): + elif args.outputfile.is_dir(): outputfolder = Path(args.outputfile) - args.outputfile = os.path.join(outputfolder, defaultoutputfile) + args.outputfile = outputfolder / defaultoutputfile else: outputfolder = Path() @@ -1516,7 +1604,7 @@ def main(args=None, argsraw=None, **kwargs): if args.brightnessattime: if args.timedays is None: print("Specify timedays") - quit() + sys.exit(1) if not args.plotviewingangle: args.plotviewingangle = [-1] if not args.colorbarcostheta and not args.colorbarphi: diff --git a/artistools/lightcurve/test_lightcurve.py b/artistools/lightcurve/test_lightcurve.py new file mode 100755 index 000000000..db11e9bc1 --- /dev/null +++ b/artistools/lightcurve/test_lightcurve.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +from pathlib import Path +from unittest import mock + +import matplotlib.axes +import numpy as np + +import artistools as at + +modelpath = at.get_config()["path_testartismodel"] +outputpath = at.get_config()["path_testoutput"] + + +@mock.patch.object(matplotlib.axes.Axes, "plot", side_effect=matplotlib.axes.Axes.plot, autospec=True) +def test_lightcurve_plot(mockplot) -> None: + at.lightcurve.plot(argsraw=[], modelpath=[modelpath], outputfile=outputpath, frompackets=False) + + arr_time_d = np.array(mockplot.call_args[0][1]) + arr_lum = np.array(mockplot.call_args[0][2]) + + assert np.isclose(arr_time_d.min(), 250.421, rtol=1e-4) + assert np.isclose(arr_time_d.max(), 349.412, rtol=1e-4) + + assert np.isclose(arr_time_d.mean(), 297.20121, rtol=1e-4) + assert np.isclose(arr_time_d.std(), 28.83886, rtol=1e-4) + + integral = np.trapz(arr_lum, arr_time_d) + assert np.isclose(integral, 2.7955e42, rtol=1e-2) + + assert np.isclose(arr_lum.mean(), 2.8993e40, rtol=1e-4) + assert np.isclose(arr_lum.std(), 1.1244e40, rtol=1e-4) + + +@mock.patch.object(matplotlib.axes.Axes, "plot", side_effect=matplotlib.axes.Axes.plot, autospec=True) +def test_lightcurve_plot_frompackets(mockplot) -> None: + at.lightcurve.plot( + argsraw=[], + modelpath=modelpath, + frompackets=True, + outputfile=Path(outputpath, "lightcurve_from_packets.pdf"), + ) + + arr_time_d = np.array(mockplot.call_args[0][1]) + arr_lum = np.array(mockplot.call_args[0][2]) + + assert np.isclose(arr_time_d.min(), 250.421, rtol=1e-4) + assert np.isclose(arr_time_d.max(), 349.412, rtol=1e-4) + + assert np.isclose(arr_time_d.mean(), 297.20121, rtol=1e-4) + assert np.isclose(arr_time_d.std(), 28.83886, rtol=1e-4) + + integral = np.trapz(arr_lum, arr_time_d) + + assert np.isclose(integral, 1.0795078026708302e41, rtol=1e-2) + + assert np.isclose(arr_lum.mean(), 1.1176e39, rtol=1e-4) + assert np.isclose(arr_lum.std(), 4.4820e38, rtol=1e-4) + + +def test_band_lightcurve_plot() -> None: + at.lightcurve.plot(argsraw=[], modelpath=modelpath, filter=["B"], outputfile=outputpath) + + +def test_band_lightcurve_subplots() -> None: + at.lightcurve.plot(argsraw=[], modelpath=modelpath, filter=["bol", "B"], outputfile=outputpath) + + +def test_colour_evolution_plot() -> None: + at.lightcurve.plot(argsraw=[], modelpath=modelpath, colour_evolution=["B-V"], outputfile=outputpath) + + +def test_colour_evolution_subplots() -> None: + at.lightcurve.plot(argsraw=[], modelpath=modelpath, colour_evolution=["U-B", "B-V"], outputfile=outputpath) diff --git a/artistools/lightcurve/viewingangleanalysis.py b/artistools/lightcurve/viewingangleanalysis.py index 71af5f3a0..2d3803794 100644 --- a/artistools/lightcurve/viewingangleanalysis.py +++ b/artistools/lightcurve/viewingangleanalysis.py @@ -1,12 +1,10 @@ -#!/usr/bin/env python3 import argparse import glob -import math import os +import sys from collections.abc import Sequence from pathlib import Path from typing import Any -from typing import Optional from typing import Union import matplotlib.pyplot as plt @@ -16,8 +14,6 @@ from matplotlib.legend_handler import HandlerTuple import artistools as at -import artistools.lightcurve - define_colours_list = [ "k", @@ -151,75 +147,50 @@ ] -def get_angle_stuff( - modelpath: Union[Path, str], args -) -> tuple[Any, Optional[np.ndarray[Any, np.dtype[Any]]], Optional[dict[int, str]]]: +def parse_directionbin_args(modelpath: Union[Path, str], args) -> tuple[Sequence[int], dict[int, str]]: modelpath = Path(modelpath) - viewing_angles = None viewing_angle_data = False if len(glob.glob(str(modelpath / "*_res.out*"))) >= 1: viewing_angle_data = True if args.plotvspecpol and os.path.isfile(modelpath / "vpkt.txt"): - angles = args.plotvspecpol + dirbins = args.plotvspecpol elif args.plotviewingangle and args.plotviewingangle[0] == -1 and viewing_angle_data: - angles = np.arange(0, 100, 1, dtype=int) + dirbins = np.arange(0, 100, 1, dtype=int) elif args.plotviewingangle and viewing_angle_data: - angles = args.plotviewingangle + dirbins = args.plotviewingangle elif ( args.calculate_costheta_phi_from_viewing_angle_numbers and args.calculate_costheta_phi_from_viewing_angle_numbers[0] == -1 ): viewing_angles = np.arange(0, 100, 1, dtype=int) - calculate_costheta_phi_for_viewing_angles(viewing_angles, modelpath) elif args.calculate_costheta_phi_from_viewing_angle_numbers: viewing_angles = args.calculate_costheta_phi_from_viewing_angle_numbers assert viewing_angles is not None - calculate_costheta_phi_for_viewing_angles(viewing_angles, modelpath) else: - angles = [None] - - angle_definition: Optional[dict] = None - if angles[0] is not None and not args.plotvspecpol: - angle_definition = calculate_costheta_phi_for_viewing_angles(angles, modelpath) - if args.average_over_phi_angle: - for dirbin in angle_definition.keys(): - assert dirbin % at.get_viewingdirection_phibincount() == 0 - costheta_label = angle_definition[dirbin].split(",")[0] - angle_definition[dirbin] = costheta_label - - if args.average_over_theta_angle: - for dirbin in angle_definition.keys(): - assert dirbin < at.get_viewingdirection_costhetabincount() - phi_label = angle_definition[dirbin].split(",")[1] - angle_definition[dirbin] = phi_label - - return angles, viewing_angles, angle_definition - - -def calculate_costheta_phi_for_viewing_angles( - dirbins: Union[np.ndarray[Any, np.dtype[Any]], Sequence[int]], modelpath: Union[Path, str, None] = None -) -> dict[int, str]: - if modelpath: - modelpath = Path(modelpath) - MABINS = at.get_viewingdirectionbincount() - if len(list(Path(modelpath).glob("*_res_00.out*"))) > 0: # if the first direction bin file exists - assert len(list(Path(modelpath).glob(f"*_res_{MABINS-1:02d}.out*"))) > 0 # check last bin exists - assert len(list(Path(modelpath).glob(f"*_res_{MABINS:02d}.out*"))) == 0 # check one beyond does not exist + dirbins = [-1] - angle_definition: dict[int, str] = {} + dirbin_definition: dict[int, str] = {} - costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_viewinganglebin_definitions() + if args.plotvspecpol: + dirbin_definition = at.get_vspec_dir_labels(modelpath=modelpath) + else: + dirbin_definition = at.get_dirbin_labels( + dirbins=dirbins, + modelpath=modelpath, + average_over_phi=args.average_over_phi_angle, + average_over_theta=args.average_over_theta_angle, + ) - for angle in dirbins: - nphibins = at.get_viewingdirection_phibincount() - costheta_index = angle // nphibins - phi_index = angle % nphibins + if args.average_over_phi_angle: + for dirbin in dirbin_definition: + assert dirbin % at.get_viewingdirection_phibincount() == 0 or dirbin == -1 - angle_definition[angle] = f"{costheta_viewing_angle_bins[costheta_index]}, {phi_viewing_angle_bins[phi_index]}" - print(f"{angle:4d} {costheta_viewing_angle_bins[costheta_index]} {phi_viewing_angle_bins[phi_index]}") + if args.average_over_theta_angle: + for dirbin in dirbin_definition: + assert dirbin < at.get_viewingdirection_costhetabincount() or dirbin == -1 - return angle_definition + return dirbins, dirbin_definition def save_viewing_angle_data_for_plotting(band_name, modelname, args): @@ -303,7 +274,7 @@ def calculate_peak_time_mag_deltam15(time, magnitude, modelname, angle, key, arg "Trying to calculate peak time / dm15 / rise time with no time range. " "This will give a stupid result. Specify args.timemin and args.timemax" ) - quit() + sys.exit(1) print( "WARNING: Both methods that can be used to fit model light curves to get " "light curve parameters (rise, decline, peak) can be impacted by how much " @@ -364,8 +335,8 @@ def match_closest_time_polyfit(reftime_polyfit): def lightcurve_polyfit(time, magnitude, args, deg=10, kernel_scale=10, lc_error=0.01): try: import george - from george import kernels import scipy.optimize as op + from george import kernels kernel = np.var(magnitude) * kernels.Matern32Kernel(kernel_scale) gp = george.GP(kernel) @@ -447,7 +418,7 @@ def make_plot_test_viewing_angle_fit( plt.axvline(x=float(time_after15days_polyfit), color="black", linestyle="--") print("time after 15 days polyfit = ", time_after15days_polyfit) plt.tight_layout() - plt.savefig(f"{key}" + "_band_" + f"{modelname}" + "_viewing_angle" + str(angle) + ".png") + plt.savefig(f"{key}_band_{modelname}_viewing_angle" + str(angle) + ".png") plt.close() @@ -458,11 +429,10 @@ def set_scatterplot_plotkwargs(modelnumber, args): plotkwargsviewingangles["alpha"] = 0.8 if args.colorbarcostheta or args.colorbarphi: update_plotkwargs_for_viewingangle_colorbar(plotkwargsviewingangles, args) + elif args.color: + plotkwargsviewingangles["color"] = args.color[modelnumber] else: - if args.color: - plotkwargsviewingangles["color"] = args.color[modelnumber] - else: - plotkwargsviewingangles["color"] = define_colours_list2[modelnumber] + plotkwargsviewingangles["color"] = define_colours_list2[modelnumber] plotkwargsangleaveraged = {} plotkwargsangleaveraged["marker"] = "o" @@ -483,15 +453,13 @@ def set_scatterplot_plotkwargs(modelnumber, args): def update_plotkwargs_for_viewingangle_colorbar( plotkwargsviewingangles: dict[str, Any], args: argparse.Namespace ) -> dict[str, Any]: - costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_viewinganglebin_definitions() + costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_costhetabin_phibin_labels() scaledmap = at.lightcurve.plotlightcurve.make_colorbar_viewingangles_colormap() - angles = np.arange(0, 100) - angle_definition = calculate_costheta_phi_for_viewing_angles(angles, args.modelpath[0]) + angles = np.arange(0, at.get_viewingdirectionbincount()) colors = [] for angle in angles: _, colorindex = at.lightcurve.plotlightcurve.get_viewinganglecolor_for_colorbar( - angle_definition, angle, costheta_viewing_angle_bins, phi_viewing_angle_bins, @@ -515,7 +483,7 @@ def set_scatterplot_plot_params(args): plt.tight_layout() if args.colorbarcostheta or args.colorbarphi: - costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_viewinganglebin_definitions() + costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_costhetabin_phibin_labels() scaledmap = at.lightcurve.plotlightcurve.make_colorbar_viewingangles_colormap() at.lightcurve.plotlightcurve.make_colorbar_viewingangles(phi_viewing_angle_bins, scaledmap, args) @@ -569,9 +537,9 @@ def make_viewing_angle_risetime_peakmag_delta_m15_scatter_plot(modelnames, key, for ii, modelname in enumerate(modelnames): viewing_angle_plot_data = pd.read_csv(key + "band_" + f"{modelname}" + "_viewing_angle_data.txt", delimiter=" ") - band_peak_mag_viewing_angles = viewing_angle_plot_data["peak_mag_polyfit"].values - band_delta_m15_viewing_angles = viewing_angle_plot_data["deltam15_polyfit"].values - band_risetime_viewing_angles = viewing_angle_plot_data["risetime_polyfit"].values + band_peak_mag_viewing_angles = viewing_angle_plot_data["peak_mag_polyfit"].to_numpy() + band_delta_m15_viewing_angles = viewing_angle_plot_data["deltam15_polyfit"].to_numpy() + band_risetime_viewing_angles = viewing_angle_plot_data["risetime_polyfit"].to_numpy() plotkwargsviewingangles, plotkwargsangleaveraged = set_scatterplot_plotkwargs(ii, args) @@ -595,10 +563,7 @@ def make_viewing_angle_risetime_peakmag_delta_m15_scatter_plot(modelnames, key, else: args.plotvalues.append((a0, a0)) if not args.noerrorbars: - if args.color: - ecolor = args.color - else: - ecolor = define_colours_list + ecolor = args.color if args.color else define_colours_list ax.errorbar( xvalues_angleaveraged, @@ -609,10 +574,7 @@ def make_viewing_angle_risetime_peakmag_delta_m15_scatter_plot(modelnames, key, capsize=2, ) - if args.label: - linelabels = args.label - else: - linelabels = modelnames + linelabels = args.label if args.label else modelnames # a0, datalabel = at.lightcurve.get_sn_sample_bol() # a0, datalabel = at.lightcurve.plot_phillips_relation_data() @@ -634,7 +596,7 @@ def make_viewing_angle_risetime_peakmag_delta_m15_scatter_plot(modelnames, key, # ax.set_xlabel(r'Decline Rate ($\Delta$m$_{15}$)', fontsize=14) if args.make_viewing_angle_peakmag_delta_m15_scatter_plot: - xlabel = r"$\Delta$m$_{15}$" + f"({key})" + xlabel = rf"$\Delta$m$_{15}$({key})" if args.make_viewing_angle_peakmag_risetime_scatter_plot: xlabel = "Rise Time [days]" @@ -666,13 +628,13 @@ def make_peak_colour_viewing_angle_plot(args): datafilename = bands[0] + "band_" + f"{modelname}" + "_viewing_angle_data.txt" viewing_angle_plot_data = pd.read_csv(datafilename, delimiter=" ") - data[f"{bands[0]}max"] = viewing_angle_plot_data["peak_mag_polyfit"].values - data[f"time_{bands[0]}max"] = viewing_angle_plot_data["risetime_polyfit"].values + data[f"{bands[0]}max"] = viewing_angle_plot_data["peak_mag_polyfit"].to_numpy() + data[f"time_{bands[0]}max"] = viewing_angle_plot_data["risetime_polyfit"].to_numpy() # Get brightness in second band at time of peak in first band if len(data[f"time_{bands[0]}max"]) != 100: print(f"All 100 angles are not in file {datafilename}. Quitting") - quit() + sys.exit(1) second_band_brightness = second_band_brightness_at_peak_first_band(data, bands, modelpath, modelnumber, args) @@ -711,7 +673,6 @@ def make_peak_colour_viewing_angle_plot(args): plt.close() -@at.diskcache(savezipped=True) def second_band_brightness_at_peak_first_band(data, bands, modelpath, modelnumber, args): second_band_brightness = [] for anglenumber, time in enumerate(data[f"time_{bands[0]}max"]): @@ -758,20 +719,16 @@ def peakmag_risetime_declinerate_init(modelpaths, filternames_conversion_dict, a modelpath = Path(modelpath) if not args.filter: - if args.plotviewingangle: - lcname = "light_curve_res.out" - else: - lcname = "light_curve.out" - lcpath = at.firstexisting(lcname, path=modelpath, tryzipped=True) - print(f"Reading {lcname}") - lightcurve_data = at.lightcurve.readfile(lcpath, modelpath, args) + lcname = "light_curve_res.out" if args.plotviewingangle else "light_curve.out" + lcpath = at.firstexisting(lcname, folder=modelpath, tryzipped=True) + lightcurve_data = at.lightcurve.readfile(lcpath) # check if doing viewing angle stuff, and if so define which data to use - angles, viewing_angles, angle_definition = get_angle_stuff(modelpath, args) + angles, angle_definition = parse_directionbin_args(modelpath, args) if not args.filter and args.plotviewingangle: lcdataframes = lightcurve_data - for index, angle in enumerate(angles): + for _index, angle in enumerate(angles): modelname = at.get_model_name(modelpath) modelnames.append(modelname) # save for later print(f"Reading spectra: {modelname}") @@ -785,7 +742,7 @@ def peakmag_risetime_declinerate_init(modelpaths, filternames_conversion_dict, a if not args.filter: plottinglist = ["lightcurve"] - for plotnumber, band_name in enumerate(plottinglist): + for _plotnumber, band_name in enumerate(plottinglist): if args.filter: time, brightness = at.lightcurve.get_band_lightcurve(lightcurve_data, band_name, args) else: @@ -837,14 +794,14 @@ def plot_viewanglebrightness_at_fixed_time(modelpath, args): nrows=1, ncols=1, sharey=True, figsize=(8, 5), tight_layout={"pad": 0.2, "w_pad": 0.0, "h_pad": 0.0} ) - angles, viewing_angles, angle_definition = at.lightcurve.get_angle_stuff(modelpath, args) + angles, angle_definition = at.lightcurve.parse_directionbin_args(modelpath, args) - costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_viewinganglebin_definitions() + costheta_viewing_angle_bins, phi_viewing_angle_bins = at.get_costhetabin_phibin_labels() scaledmap = at.lightcurve.plotlightcurve.make_colorbar_viewingangles_colormap() plotkwargs = {} - lcdataframes = at.lightcurve.readfile(modelpath / "light_curve_res.out", args) + lcdataframes = at.lightcurve.readfile(modelpath / "light_curve_res.out") timetoplot = at.match_closest_time(reftime=args.timedays, searchtimes=lcdataframes[0]["time"]) print(timetoplot) @@ -852,7 +809,7 @@ def plot_viewanglebrightness_at_fixed_time(modelpath, args): for angleindex, lcdata in enumerate(lcdataframes): angle = angleindex plotkwargs, _ = at.lightcurve.plotlightcurve.get_viewinganglecolor_for_colorbar( - angle_definition, angle, costheta_viewing_angle_bins, phi_viewing_angle_bins, scaledmap, plotkwargs, args + angle, costheta_viewing_angle_bins, phi_viewing_angle_bins, scaledmap, plotkwargs, args ) rowattime = lcdata.loc[lcdata["time"] == float(timetoplot)] diff --git a/artistools/lightcurve/writebollightcurvedata.py b/artistools/lightcurve/writebollightcurvedata.py index 51da35e5d..b7ee74cf5 100644 --- a/artistools/lightcurve/writebollightcurvedata.py +++ b/artistools/lightcurve/writebollightcurvedata.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 from pathlib import Path import numpy as np @@ -6,13 +5,12 @@ from astropy import units as u import artistools as at -import artistools.spectra def get_bol_lc_from_spec(modelpath): res_specdata = at.spectra.read_spec_res(modelpath) # print(res_specdata) - timearray = res_specdata[0].columns.values[1:] + timearray = res_specdata[0].columns.to_numpy()[1:] times = [time for time in timearray if 5 < float(time) < 80] lightcurvedata = {"time": times} @@ -21,9 +19,9 @@ def get_bol_lc_from_spec(modelpath): for timestep, time in enumerate(timearray): time = float(time) if 5 < time < 80: - spectrum = at.spectra.get_res_spectrum( - modelpath, timestep, timestep, angle=angle, res_specdata=res_specdata - ) + spectrum = at.spectra.get_spectrum( + modelpath=modelpath, directionbins=[angle], timestepmin=timestep, timestepmax=timestep + )[angle] integrated_flux = np.trapz(spectrum["f_lambda"], spectrum["lambda_angstroms"]) integrated_luminosity = integrated_flux * 4 * np.pi * np.power(u.Mpc.to("cm"), 2) bol_luminosity.append(integrated_luminosity) @@ -38,12 +36,9 @@ def get_bol_lc_from_spec(modelpath): def get_bol_lc_from_lightcurveout(modelpath: Path, res: bool = False) -> pd.DataFrame: - if res: - lcfilename = "light_curve_res.out" - else: - lcfilename = "light_curve.out" + lcfilename = "light_curve_res.out" if res else "light_curve.out" lcdata = pd.read_csv(modelpath / lcfilename, delim_whitespace=True, header=None, names=["time", "lum", "lum_cmf"]) - lcdataframes = at.gather_res_data(lcdata, index_of_repeated_value=0) + lcdataframes = at.split_dataframe_dirbins(lcdata, index_of_repeated_value=0) times = lcdataframes[0]["time"] lightcurvedata = {"time": times} diff --git a/artistools/linefluxes.py b/artistools/linefluxes.py old mode 100644 new mode 100755 index ed68f348c..04ab10f7c --- a/artistools/linefluxes.py +++ b/artistools/linefluxes.py @@ -15,26 +15,20 @@ from astropy import units as u import artistools as at -import artistools.packets - -# from functools import lru_cache -# import matplotlib.ticker as ticker -# from astropy import constants as const def get_packets_with_emtype_onefile(emtypecolumn, lineindices, packetsfile): import gzip try: - dfpackets = at.packets.readfile(packetsfile, type="TYPE_ESCAPE", escape_type="TYPE_RPKT") - except gzip.BadGzipFile: + dfpackets = at.packets.readfile(packetsfile, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT") + except gzip.BadGzipFile as exc: print(f"Bad file: {packetsfile}") - raise gzip.BadGzipFile + raise gzip.BadGzipFile from exc return dfpackets.query(f"{emtypecolumn} in @lineindices", inplace=False).copy() -@at.diskcache(savezipped=True) def get_packets_with_emtype(modelpath, emtypecolumn, lineindices, maxpacketfiles=None): packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles=maxpacketfiles) nprocs_read = len(packetsfiles) @@ -70,7 +64,7 @@ def calculate_timebinned_packet_sum(dfpackets, timearrayplusend): def get_line_fluxes_from_packets( emtypecolumn, emfeatures, modelpath, maxpacketfiles=None, arr_tstart=None, arr_tend=None -): +) -> pd.DataFrame: if arr_tstart is None: arr_tstart = at.get_timestep_times_float(modelpath, loc="start") if arr_tend is None: @@ -111,9 +105,7 @@ def get_line_fluxes_from_packets( return lcdata -def get_line_fluxes_from_pops(emtypecolumn, emfeatures, modelpath, arr_tstart=None, arr_tend=None): - import artistools.nltepops - +def get_line_fluxes_from_pops(emfeatures, modelpath, arr_tstart=None, arr_tend=None) -> pd.DataFrame: if arr_tstart is None: arr_tstart = at.get_timestep_times_float(modelpath, loc="start") if arr_tend is None: @@ -144,8 +136,8 @@ def get_line_fluxes_from_pops(emtypecolumn, emfeatures, modelpath, arr_tstart=No ion = adata.query("Z == @feature.atomic_number and ion_stage == @feature.ion_stage").iloc[0] for timeindex, timedays in enumerate(arr_tmid): - v_inner = modeldata.velocity_inner.values * u.km / u.s - v_outer = modeldata.velocity_outer.values * u.km / u.s + v_inner = modeldata.velocity_inner.to_numpy() * u.km / u.s + v_outer = modeldata.velocity_outer.to_numpy() * u.km / u.s t_sec = timedays * u.day shell_volumes = ((4 * math.pi / 3) * ((v_outer * t_sec) ** 3 - (v_inner * t_sec) ** 3)).to("cm3").value @@ -176,7 +168,7 @@ def get_line_fluxes_from_pops(emtypecolumn, emfeatures, modelpath, arr_tstart=No ) * u.eV.to("erg") # l = delta_ergs * A_val * levelpop * (shell_volumes[modelgridindex] + unaccounted_shellvol) - # print(f' {modelgridindex} outer_velocity {modeldata.velocity_outer.values[modelgridindex]}' + # print(f' {modelgridindex} outer_velocity {modeldata.velocity_outer.to_numpy()[modelgridindex]}' # f' km/s shell_vol: {shell_volumes[modelgridindex] + unaccounted_shellvol} cm3' # f' n_level {levelpop} cm-3 shell_Lum {l} erg/s') @@ -212,18 +204,18 @@ def get_closelines( dflinelist = at.get_linelist_dataframe(modelpath) dflinelistclosematches = dflinelist.query("atomic_number == @atomic_number and ionstage == @ion_stage").copy() if lambdamin > 0: - dflinelistclosematches.query("@lambdamin < lambda_angstroms", inplace=True) + dflinelistclosematches = dflinelistclosematches.query("@lambdamin < lambda_angstroms") if lambdamax > 0: - dflinelistclosematches.query("@lambdamax > lambda_angstroms", inplace=True) + dflinelistclosematches = dflinelistclosematches.query("@lambdamax > lambda_angstroms") if lowerlevelindex >= 0: - dflinelistclosematches.query("lowerlevelindex==@lowerlevelindex", inplace=True) + dflinelistclosematches = dflinelistclosematches.query("lowerlevelindex==@lowerlevelindex") if upperlevelindex >= 0: - dflinelistclosematches.query("upperlevelindex==@upperlevelindex", inplace=True) + dflinelistclosematches = dflinelistclosematches.query("upperlevelindex==@upperlevelindex") # print(dflinelistclosematches) - linelistindices = tuple(dflinelistclosematches.index.values) - upperlevelindicies = tuple(dflinelistclosematches.upperlevelindex.values) - lowerlevelindicies = tuple(dflinelistclosematches.lowerlevelindex.values) + linelistindices = tuple(dflinelistclosematches.index.to_numpy()) + upperlevelindicies = tuple(dflinelistclosematches.upperlevelindex.to_numpy()) + lowerlevelindicies = tuple(dflinelistclosematches.lowerlevelindex.to_numpy()) lowestlambda = dflinelistclosematches.lambda_angstroms.min() highestlamba = dflinelistclosematches.lambda_angstroms.max() colname = f"flux_{at.get_ionstring(atomic_number, ion_stage, nospace=True)}_{approxlambda}" @@ -302,14 +294,14 @@ def make_flux_ratio_plot(args): pd.set_option("display.width", 150) pd.options.display.max_rows = 500 - for seriesindex, (modelpath, modellabel, modelcolor) in enumerate(zip(args.modelpath, args.label, args.color)): + for _seriesindex, (modelpath, modellabel, modelcolor) in enumerate(zip(args.modelpath, args.label, args.color)): print(f"====> {modellabel}") emfeatures = get_labelandlineindices(modelpath, tuple(args.emfeaturesearch)) if args.frompops: dflcdata = get_line_fluxes_from_pops( - args.emtypecolumn, emfeatures, modelpath, arr_tstart=args.timebins_tstart, arr_tend=args.timebins_tend + emfeatures, modelpath, arr_tstart=args.timebins_tstart, arr_tend=args.timebins_tend ) else: dflcdata = get_line_fluxes_from_packets( @@ -321,7 +313,7 @@ def make_flux_ratio_plot(args): arr_tend=args.timebins_tend, ) - dflcdata.eval(f"fratio = {emfeatures[1].colname} / {emfeatures[0].colname}", inplace=True) + dflcdata = dflcdata.eval(f"fratio = {emfeatures[1].colname} / {emfeatures[0].colname}") axis.set_ylabel( r"F$_{\mathrm{" + emfeatures[1].featurelabel + r"}}$ / F$_{\mathrm{" + emfeatures[0].featurelabel + r"}}$" ) @@ -332,7 +324,7 @@ def make_flux_ratio_plot(args): axis.plot( dflcdata.time, - dflcdata["fratio"], + dflcdata.fratio, label=modellabel, marker="x", lw=0, @@ -360,7 +352,7 @@ def make_flux_ratio_plot(args): ) amodels = {} - for index, row in femis.iterrows(): + for _index, row in femis.iterrows(): modelname = row.file.replace("fig-nne_Te_allcells-", "").replace(f"-{row.epoch}d.txt", "") if modelname not in amodels: amodels[modelname] = ([], []) @@ -422,7 +414,6 @@ def make_flux_ratio_plot(args): plt.close() -@at.diskcache(savezipped=True) def get_packets_with_emission_conditions(modelpath, emtypecolumn, lineindices, tstart, tend, maxpacketfiles=None): estimators = at.estimators.read_estimators(modelpath, get_ion_values=False, get_heatingcooling=False) @@ -434,7 +425,7 @@ def get_packets_with_emission_conditions(modelpath, emtypecolumn, lineindices, t # model_tmids = at.get_timestep_times_float(modelpath, loc='mid') # arr_velocity_mid = tuple(list([(float(v1) + float(v2)) * 0.5 for v1, v2 in zip( - # modeldata['velocity_inner'].values, modeldata['velocity_outer'].values)])) + # modeldata['velocity_inner'].to_numpy(), modeldata['velocity_outer'].to_numpy())])) # from scipy.interpolate import interp1d # interp_log10nne, interp_te = {}, {} @@ -516,17 +507,17 @@ def plot_nne_te_points(axis, serieslabel, em_log10nne, em_Te, normtotalpackets, def plot_nne_te_bars(axis, serieslabel, em_log10nne, em_Te, color): if len(em_log10nne) == 0: return - errorbarkwargs = dict( - xerr=np.std(em_log10nne), - yerr=np.std(em_Te), - color="black", - markersize=10.0, - fillstyle="full", - capthick=4, - capsize=15, - linewidth=4.0, - alpha=1.0, - ) + errorbarkwargs = { + "xerr": np.std(em_log10nne), + "yerr": np.std(em_Te), + "color": "black", + "markersize": 10.0, + "fillstyle": "full", + "capthick": 4, + "capsize": 15, + "linewidth": 4.0, + "alpha": 1.0, + } # black larger one for an outline axis.errorbar(np.mean(em_log10nne), np.mean(em_Te), **errorbarkwargs) errorbarkwargs["markersize"] -= 2 @@ -539,8 +530,6 @@ def plot_nne_te_bars(axis, serieslabel, em_log10nne, em_Te, color): def make_emitting_regions_plot(args): - import artistools.estimators - # font = {'size': 16} # matplotlib.rc('font', **font) # 'floers_te_nne.json', @@ -559,9 +548,7 @@ def make_emitting_regions_plot(args): floers_te_nne = json.loads(data_file.read()) # give an ordering and index to dict items - refdatakeys[refdataindex] = [ - t for t in sorted(floers_te_nne.keys(), key=lambda x: float(x)) - ] # strings, not floats + refdatakeys[refdataindex] = sorted(floers_te_nne.keys(), key=float) # strings, not floats refdatatimes[refdataindex] = np.array([float(t) for t in refdatakeys[refdataindex]]) refdatapoints[refdataindex] = [floers_te_nne[t] for t in refdatakeys[refdataindex]] print(f"{refdatafilename} data available for times: {list(refdatatimes[refdataindex])}") @@ -619,8 +606,8 @@ def make_emitting_regions_plot(args): emdata_all[modelindex][(tmid, feature.colname)] = {"em_log10nne": [], "em_Te": []} else: emdata_all[modelindex][(tmid, feature.colname)] = { - "em_log10nne": dfpackets_selected.em_log10nne.values, - "em_Te": dfpackets_selected.em_Te.values, + "em_log10nne": dfpackets_selected.em_log10nne.to_numpy(), + "em_Te": dfpackets_selected.em_Te.to_numpy(), } estimators = at.estimators.read_estimators(modelpath, get_ion_values=False, get_heatingcooling=False) @@ -662,7 +649,7 @@ def make_emitting_regions_plot(args): tight_layout={"pad": 0.2, "w_pad": 0.0, "h_pad": 0.2}, ) - for refdataindex, f in enumerate(refdatafilenames): + for refdataindex, _f in enumerate(refdatafilenames): timeindex = np.abs(refdatatimes[refdataindex] - tmid).argmin() axis.plot( refdatapoints[refdataindex][timeindex]["ne"], @@ -816,10 +803,6 @@ def addargs(parser: argparse.ArgumentParser) -> None: parser.add_argument("-emfeaturesearch", default=[], nargs="*", help="List of tuples (TODO explain)") - # parser.add_argument('-emtypecolumn', default='trueemissiontype', choices=['emissiontype', 'trueemissiontype'], - # help='Packet property for emission type - first thermal emission (trueemissiontype) ' - # 'or last emission type (emissiontype)') - parser.add_argument( "--frompops", action="store_true", help="Sum up internal emissivity instead of outgoing packets" ) diff --git a/artistools/logfiles.py b/artistools/logfiles.py old mode 100644 new mode 100755 index 841da71c5..b1de953e2 --- a/artistools/logfiles.py +++ b/artistools/logfiles.py @@ -1,12 +1,9 @@ #!/usr/bin/env python3 import argparse -import glob import multiprocessing -import sys from pathlib import Path import matplotlib.pyplot as plt -import pandas as pd import artistools as at @@ -45,7 +42,7 @@ def read_time_taken(logfilepaths): for logfilepath in logfilepaths: mpi_process = int(str(logfilepath).split("/")[-1].split("-")[0].split("_")[-1]) - with open(logfilepath, "r") as logfile: + with open(logfilepath) as logfile: lineswithtimes = [line.split(" ") for line in logfile if "took" in line] # for line in lineswithtimes: diff --git a/artistools/macroatom.py b/artistools/macroatom.py old mode 100644 new mode 100755 index 3c6116af2..838baaf05 --- a/artistools/macroatom.py +++ b/artistools/macroatom.py @@ -7,13 +7,9 @@ import matplotlib.pyplot as plt import pandas as pd -from astropy import constants as const import artistools as at -# from astropy import units as u -# import matplotlib.ticker as ticker - defaultoutputfile = "plotmacroatom_cell{0:03d}_{1:03d}-{2:03d}.pdf" @@ -46,14 +42,11 @@ def main(args=None, argsraw=None, **kwargs): atomic_number = at.get_atomic_number(args.element.lower()) if atomic_number < 1: print(f"Could not find element '{args.element}'") - return + raise AssertionError timestepmin = args.timestep - if not args.timestepmax or args.timestepmax < 0: - timestepmax = timestepmin - else: - timestepmax = args.timestepmax + timestepmax = timestepmin if not args.timestepmax or args.timestepmax < 0 else args.timestepmax input_files = glob.glob(os.path.join(args.modelpath, "macroatom_????.out*"), recursive=True) + glob.glob( os.path.join(args.modelpath, "*/macroatom_????.out*"), recursive=True @@ -61,9 +54,9 @@ def main(args=None, argsraw=None, **kwargs): if not input_files: print("No macroatom files found") - return 1 - else: - dfall = read_files(input_files, args.modelgridindex, timestepmin, timestepmax, atomic_number) + raise FileNotFoundError + + dfall = read_files(input_files, args.modelgridindex, timestepmin, timestepmax, atomic_number) specfilename = os.path.join(args.modelpath, "spec.out") @@ -72,7 +65,7 @@ def main(args=None, argsraw=None, **kwargs): if not os.path.isfile(specfilename): print(f"Could not find {specfilename}") - return 1 + raise FileNotFoundError outputfile = args.outputfile.format(args.modelgridindex, timestepmin, timestepmax) make_plot( @@ -121,8 +114,8 @@ def make_plot( fontsize=8, ) - lambda_cmf_in = 2.99792458e18 / dfmacroatom["nu_cmf_in"].values - lambda_cmf_out = 2.99792458e18 / dfmacroatom["nu_cmf_out"].values + lambda_cmf_in = 2.99792458e18 / dfmacroatom["nu_cmf_in"].to_numpy() + lambda_cmf_out = 2.99792458e18 / dfmacroatom["nu_cmf_out"].to_numpy() # axis.scatter(lambda_cmf_in, lambda_cmf_out, s=1, alpha=0.5, edgecolor='none') axis.plot( lambda_cmf_in, @@ -157,20 +150,16 @@ def read_files(files, modelgridindex=None, timestepmin=None, timestepmax=None, a df_thisfile = pd.read_csv(filepath, delim_whitespace=True) # df_thisfile[['modelgridindex', 'timestep']].apply(pd.to_numeric) if modelgridindex: - df_thisfile.query("modelgridindex==@modelgridindex", inplace=True) + df_thisfile = df_thisfile.query("modelgridindex==@modelgridindex") if timestepmin is not None: - df_thisfile.query("timestep>=@timestepmin", inplace=True) + df_thisfile = df_thisfile.query("timestep>=@timestepmin") if timestepmax: - df_thisfile.query("timestep<=@timestepmax", inplace=True) + df_thisfile = df_thisfile.query("timestep<=@timestepmax") if atomic_number: - df_thisfile.query("Z==@atomic_number", inplace=True) - - if df_thisfile is not None: - if len(df_thisfile) > 0: - if dfall is None: - dfall = df_thisfile.copy() - else: - dfall = dfall.append(df_thisfile.copy(), ignore_index=True) + df_thisfile = df_thisfile.query("Z==@atomic_number") + + if df_thisfile is not None and len(df_thisfile) > 0: + dfall = df_thisfile.copy() if dfall is None else dfall.append(df_thisfile.copy(), ignore_index=True) if dfall is None or len(dfall) == 0: print("No data found") diff --git a/artistools/misc.py b/artistools/misc.py index c4900226c..f25a1f55c 100644 --- a/artistools/misc.py +++ b/artistools/misc.py @@ -1,12 +1,8 @@ -#!/usr/bin/env python3 import argparse import gzip import io -import lzma import math import os -import sys -import time from collections import namedtuple from collections.abc import Collection from collections.abc import Iterable @@ -21,16 +17,15 @@ from typing import Optional from typing import Union +import lz4.frame import numpy as np import pandas as pd +import polars as pl +import pyzstd +import xz import artistools as at -# import inspect -# from functools import partial -# import scipy.signal - - roman_numerals = ( "", "I", @@ -59,18 +54,18 @@ class CustomArgHelpFormatter(argparse.ArgumentDefaultsHelpFormatter): def add_arguments(self, actions: Iterable[argparse.Action]) -> None: def my_sort(arg: Any) -> str: - opstr = arg.option_strings[0] if len(arg.option_strings) > 0 else "" + opstr: str = arg.option_strings[0] if len(arg.option_strings) > 0 else "" # chars = 'abcdefghijklmnopqrstuvwxyz-' opstr = opstr.upper().replace("-", "z") # push dash chars below alphabet return opstr actions = sorted(actions, key=my_sort) - super(CustomArgHelpFormatter, self).add_arguments(actions) + super().add_arguments(actions) class AppendPath(argparse.Action): - def __call__(self, parser, args, values, option_string=None) -> None: # type: ignore + def __call__(self, parser, args, values, option_string=None) -> None: # type: ignore[no-untyped-def] # noqa: ARG002 # if getattr(args, self.dest) is None: # setattr(args, self.dest, []) if isinstance(values, Iterable): @@ -86,15 +81,7 @@ def __call__(self, parser, args, values, option_string=None) -> None: # type: i setattr(args, self.dest, Path(values)) -def make_namedtuple(typename: str, **fields: dict) -> tuple[Any, ...]: - """Make a namedtuple from a dictionary of attributes and values. - Example: make_namedtuple('mytuple', x=2, y=3)""" - return namedtuple(typename, fields)(*fields.values()) - - -def showtimesteptimes( - modelpath: Optional[Path] = None, numberofcolumns: int = 5, args: Optional[argparse.Namespace] = None -) -> None: +def showtimesteptimes(modelpath: Optional[Path] = None, numberofcolumns: int = 5) -> None: """Print a table showing the timesteps and their corresponding times.""" if modelpath is None: modelpath = Path() @@ -123,7 +110,7 @@ def get_composition_data(filename: Union[Path, str]) -> pd.DataFrame: columns = ("Z,nions,lowermost_ionstage,uppermost_ionstage,nlevelsmax_readin,abundance,mass,startindex").split(",") rowdfs = [] - with open(filename, "r") as fcompdata: + with open(filename) as fcompdata: nelements = int(fcompdata.readline()) fcompdata.readline() # T_preset fcompdata.readline() # homogeneous_abundances @@ -131,11 +118,11 @@ def get_composition_data(filename: Union[Path, str]) -> pd.DataFrame: for _ in range(nelements): line = fcompdata.readline() linesplit = line.split() - row_list = list([int(x) for x in linesplit[:5]]) + list([float(x) for x in linesplit[5:]]) + [startindex] + row_list = [int(x) for x in linesplit[:5]] + [float(x) for x in linesplit[5:]] + [startindex] rowdfs.append(pd.DataFrame([row_list], columns=columns)) - startindex += int(rowdfs[-1]["nions"]) + startindex += int(rowdfs[-1].iloc[0]["nions"]) compdf = pd.concat(rowdfs, ignore_index=True) @@ -146,7 +133,7 @@ def get_composition_data_from_outputfile(modelpath: Path) -> pd.DataFrame: """Read ion list from output file""" atomic_composition = {} - output = open(modelpath / "output_0-0.txt", "rt").read().splitlines() + output = open(modelpath / "output_0-0.txt").read().splitlines() Z: Optional[int] = None ioncount = 0 for row in output: @@ -160,31 +147,82 @@ def get_composition_data_from_outputfile(modelpath: Path) -> pd.DataFrame: assert Z is not None atomic_composition[Z] = ioncount - composition_df = pd.DataFrame( - [(Z, atomic_composition[Z]) for Z in atomic_composition.keys()], columns=["Z", "nions"] - ) + composition_df = pd.DataFrame([(Z, atomic_composition[Z]) for Z in atomic_composition], columns=["Z", "nions"]) composition_df["lowermost_ionstage"] = [1] * composition_df.shape[0] composition_df["uppermost_ionstage"] = composition_df["nions"] return composition_df -def gather_res_data(res_df: pd.DataFrame, index_of_repeated_value: int = 1) -> dict[int, pd.DataFrame]: +def split_dataframe_dirbins( + res_df: Union[pl.DataFrame, pd.DataFrame], index_of_repeated_value: int = 0, output_polarsdf: bool = False +) -> dict[int, Union[pd.DataFrame, pl.DataFrame]]: """res files repeat output for each angle. - index_of_repeated_value is the value to look for repeating eg. time of ts 0. - In spec_res files it's 1, but in lc_res file it's 0""" - index_to_split = res_df.index[res_df.iloc[:, index_of_repeated_value] == res_df.iloc[0, index_of_repeated_value]] - res_data = {} - for i, index_value in enumerate(index_to_split): - if index_value != index_to_split[-1]: - chunk = res_df.iloc[index_to_split[i] : index_to_split[i + 1], :] - else: - chunk = res_df.iloc[index_to_split[i] :, :] - res_data[i] = chunk + index_of_repeated_value is the column index to look for repeating eg. time of ts 0. + In spec_res files it's 1 , but in lc_res file it's 0""" + + if isinstance(res_df, pd.DataFrame): + res_df = pl.from_pandas(res_df) + + indexes_to_split = pl.arg_where( + res_df[:, index_of_repeated_value] == res_df[0, index_of_repeated_value], eager=True + ) + + res_data: dict[int, pl.DataFrame] = {} + prev_dfshape = None + for i, index_value in enumerate(indexes_to_split): + chunk = ( + res_df[indexes_to_split[i] : indexes_to_split[i + 1], :] + if index_value != indexes_to_split[-1] + else res_df[indexes_to_split[i] :, :] + ) - assert len(res_data) == 100 + # the number of timesteps should match for all direction bins + assert prev_dfshape is None or prev_dfshape == chunk.shape + prev_dfshape = chunk.shape + + res_data[i] = chunk if output_polarsdf else chunk.to_pandas(use_pyarrow_extension_array=True) + + assert len(res_data) == at.get_viewingdirectionbincount() return res_data +def average_direction_bins( + dirbindataframes: dict[int, pl.DataFrame], + overangle: Literal["phi", "theta"], +) -> dict[int, pl.DataFrame]: + """Will average dict of direction-binned polars DataFrames according to the phi or theta angle""" + dirbincount = at.get_viewingdirectionbincount() + nphibins = at.get_viewingdirection_phibincount() + + assert overangle in ["phi", "theta"] + if overangle == "phi": + start_bin_range = range(0, dirbincount, nphibins) + elif overangle == "theta": + ncosthetabins = at.get_viewingdirection_costhetabincount() + start_bin_range = range(0, nphibins) + + # we will make a copy to ensure that we don't cause side effects from altering the original DataFrames + # that might be returned again later by an lru_cached function + dirbindataframesout: dict[int, pd.DataFrame] = {} + + for start_bin in start_bin_range: + dirbindataframesout[start_bin] = dirbindataframes[start_bin] + + contribbins = ( + range(start_bin + 1, start_bin + nphibins) + if overangle == "phi" + else range(start_bin + ncosthetabins, dirbincount, ncosthetabins) + ) + + for dirbin_contrib in contribbins: + dirbindataframesout[start_bin] += dirbindataframes[dirbin_contrib] + + dirbindataframesout[start_bin] /= 1 + len(contribbins) # every nth bin is the average of n bins + print(f"bin number {start_bin} = the average of bins {[start_bin, *list(contribbins)]}") + + return dirbindataframesout + + def match_closest_time(reftime: float, searchtimes: list[Any]) -> str: """Get time closest to reftime in list of times (searchtimes)""" return str("{}".format(min([float(x) for x in searchtimes], key=lambda x: abs(x - reftime)))) @@ -193,7 +231,7 @@ def match_closest_time(reftime: float, searchtimes: list[Any]) -> str: def get_vpkt_config(modelpath: Union[Path, str]) -> dict[str, Any]: filename = Path(modelpath, "vpkt.txt") vpkt_config: dict[str, Any] = {} - with open(filename, "r") as vpkt_txt: + with open(filename) as vpkt_txt: vpkt_config["nobsdirections"] = int(vpkt_txt.readline()) vpkt_config["cos_theta"] = [float(x) for x in vpkt_txt.readline().split()] vpkt_config["phi"] = [float(x) for x in vpkt_txt.readline().split()] @@ -201,14 +239,14 @@ def get_vpkt_config(modelpath: Union[Path, str]) -> dict[str, Any]: if nspecflag == 1: vpkt_config["nspectraperobs"] = int(vpkt_txt.readline()) - for i in range(vpkt_config["nspectraperobs"]): + for _ in range(vpkt_config["nspectraperobs"]): vpkt_txt.readline() else: vpkt_config["nspectraperobs"] = 1 - vpkt_config["time_limits_enabled"], vpkt_config["initial_time"], vpkt_config["final_time"] = [ + vpkt_config["time_limits_enabled"], vpkt_config["initial_time"], vpkt_config["final_time"] = ( int(x) for x in vpkt_txt.readline().split() - ] + ) return vpkt_config @@ -219,14 +257,11 @@ def get_grid_mapping(modelpath: Union[Path, str]) -> tuple[dict[int, list[int]], a dict with the associated model grid cell of each propagration cell.""" modelpath = Path(modelpath) - if modelpath.is_dir(): - filename = firstexisting("grid.out", tryzipped=True, path=modelpath) - else: - filename = Path(modelpath) + filename = firstexisting("grid.out", tryzipped=True, folder=modelpath) if modelpath.is_dir() else Path(modelpath) assoc_cells: dict[int, list[int]] = {} mgi_of_propcells: dict[int, int] = {} - with open(filename, "r") as fgrid: + with open(filename) as fgrid: for line in fgrid: row = line.split() propcellid, mgi = int(row[0]), int(row[1]) @@ -244,7 +279,7 @@ def get_wid_init_at_tmin(modelpath: Path) -> float: tmin = get_timestep_times_float(modelpath, loc="start")[0] * day_to_sec _, modelmeta = at.get_modeldata(modelpath) - rmax = modelmeta["vmax_cmps"] * tmin + rmax: float = modelmeta["vmax_cmps"] * tmin coordmax0 = rmax ncoordgrid0 = 50 @@ -262,36 +297,39 @@ 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 - dfmodel, modelmeta = at.get_modeldata(modelpath, getheadersonly=True, skipabundancecolumns=True) + dfmodel, modelmeta = at.get_modeldata(modelpath, getheadersonly=True) assert modelmeta["dimensions"] == 3 ngridpoints = len(dfmodel) - xmax = modelmeta["vmax_cmps"] * modelmeta["t_model_init_days"] * (24 * 60 * 60) - - ncoordgridx = round(ngridpoints ** (1.0 / 3.0)) + xmax = modelmeta["vmax_cmps"] * modelmeta["t_model_init_days"] * 86400 + ncoordgridx: int = round(ngridpoints ** (1.0 / 3.0)) assert xmax is not None - wid_init = 2 * xmax / ncoordgridx + wid_init = 2.0 * xmax / ncoordgridx return wid_init def get_syn_dir(modelpath: Path) -> Sequence[int]: - with open(modelpath / "syn_dir.txt", "rt") as syn_dir_file: + if not (modelpath / "syn_dir.txt").is_file(): + print(f"{modelpath / 'syn_dir.txt'} does not exist. using x,y,z = 0,0,1") + return (0, 0, 1) + + with open(modelpath / "syn_dir.txt") as syn_dir_file: syn_dir = [int(x) for x in syn_dir_file.readline().split()] return syn_dir -def vec_len(vec: Union[Sequence, np.ndarray[Any, Any]]) -> float: - return np.sqrt(np.dot(vec, vec)) +def vec_len(vec: Union[Sequence[float], np.ndarray[Any, np.dtype[np.float64]]]) -> float: + return float(np.sqrt(np.dot(vec, vec))) @lru_cache(maxsize=16) def get_nu_grid(modelpath: Path) -> np.ndarray[Any, np.dtype[np.float64]]: """Get an array of frequencies at which the ARTIS spectra are binned by exspec.""" - specfilename = firstexisting(["spec.out", "specpol.out"], path=modelpath, tryzipped=True) + specfilename = firstexisting(["spec.out", "specpol.out"], folder=modelpath, tryzipped=True) specdata = pd.read_csv(specfilename, delim_whitespace=True) - return specdata.loc[:, "0"].values + return specdata.loc[:, "0"].to_numpy() def get_deposition(modelpath: Path) -> pd.DataFrame: @@ -303,7 +341,7 @@ def get_deposition(modelpath: Path) -> pd.DataFrame: ts_mids = get_timestep_times_float(modelpath, loc="mid") - with open(depfilepath, "r") as fdep: + with open(depfilepath) as fdep: filepos = fdep.tell() line = fdep.readline() if line.startswith("#"): @@ -324,17 +362,6 @@ def get_deposition(modelpath: Path) -> pd.DataFrame: return depdata -@lru_cache(maxsize=16) -def get_timestep_times(modelpath: Path) -> list[str]: - """Return a list of the mid time in days of each timestep from a spec.out file.""" - try: - specfilename = firstexisting(["spec.out", "specpol.out"], path=modelpath, tryzipped=True) - time_columns = pd.read_csv(specfilename, delim_whitespace=True, nrows=0) - return list(time_columns.columns[1:]) - except FileNotFoundError: - return [f"{tdays:.3f}" for tdays in get_timestep_times_float(modelpath, loc="mid")] - - @lru_cache(maxsize=16) def get_timestep_times_float( modelpath: Union[Path, str], loc: Literal["mid", "start", "end", "delta"] = "mid" @@ -355,15 +382,15 @@ def get_timestep_times_float( if tsfilepath.exists(): dftimesteps = pd.read_csv(tsfilepath, delim_whitespace=True, escapechar="#", index_col="timestep") if loc == "mid": - return dftimesteps.tmid_days.values - elif loc == "start": - return dftimesteps.tstart_days.values - elif loc == "end": - return dftimesteps.tstart_days.values + dftimesteps.twidth_days.values - elif loc == "delta": - return dftimesteps.twidth_days.values - else: - raise ValueError("loc must be one of 'mid', 'start', 'end', or 'delta'") + return dftimesteps.tmid_days.to_numpy() + if loc == "start": + return dftimesteps.tstart_days.to_numpy() + if loc == "end": + return dftimesteps.tstart_days.to_numpy() + dftimesteps.twidth_days.to_numpy() + if loc == "delta": + return dftimesteps.twidth_days.to_numpy() + + raise ValueError("loc must be one of 'mid', 'start', 'end', or 'delta'") # older versions of Artis always used logarithmic timesteps and didn't produce a timesteps.out file inputparams = get_inputparams(modelpath) @@ -373,17 +400,17 @@ def get_timestep_times_float( if loc == "mid": tmids = np.array([tmin * math.exp((ts + 0.5) * dlogt) for ts in timesteps]) return tmids - elif loc == "start": + if loc == "start": tstarts = np.array([tmin * math.exp(ts * dlogt) for ts in timesteps]) return tstarts - elif loc == "end": + if loc == "end": tends = np.array([tmin * math.exp((ts + 1) * dlogt) for ts in timesteps]) return tends - elif loc == "delta": + if loc == "delta": tdeltas = np.array([tmin * (math.exp((ts + 1) * dlogt) - math.exp(ts * dlogt)) for ts in timesteps]) return tdeltas - else: - raise ValueError("loc must be one of 'mid', 'start', 'end', or 'delta'") + + raise ValueError("loc must be one of 'mid', 'start', 'end', or 'delta'") def get_timestep_of_timedays(modelpath: Path, timedays: Union[str, float]) -> int: @@ -405,8 +432,6 @@ def get_timestep_of_timedays(modelpath: Path, timedays: Union[str, float]) -> in return ts raise ValueError(f"Could not find timestep bracketing time {timedays_float}") - assert False - return def get_time_range( @@ -414,7 +439,7 @@ def get_time_range( timestep_range_str: Optional[str] = None, timemin: Optional[float] = None, timemax: Optional[float] = None, - timedays_range_str: Optional[str] = None, + timedays_range_str: Union[None, str, float] = None, ) -> tuple[int, int, Optional[float], Optional[float]]: """Handle a time range specified in either days or timesteps.""" # assertions make sure time is specified either by timesteps or times in days, but not both! @@ -425,15 +450,15 @@ def get_time_range( if timemin and timemin > tends[-1]: print(f"{get_model_name(modelpath)}: WARNING timemin {timemin} is after the last timestep at {tends[-1]:.1f}") return -1, -1, timemin, timemax - elif timemax and timemax < tstarts[0]: + if timemax and timemax < tstarts[0]: print( f"{get_model_name(modelpath)}: WARNING timemax {timemax} is before the first timestep at {tstarts[0]:.1f}" ) return -1, -1, timemin, timemax if timestep_range_str is not None: - if isinstance(timestep_range_str, str) and "-" in timestep_range_str: - timestepmin, timestepmax = [int(nts) for nts in timestep_range_str.split("-")] + if "-" in timestep_range_str: + timestepmin, timestepmax = (int(nts) for nts in timestep_range_str.split("-")) else: timestepmin = int(timestep_range_str) timestepmax = timestepmin @@ -443,7 +468,7 @@ def get_time_range( timestepmax = None if timedays_range_str is not None: if isinstance(timedays_range_str, str) and "-" in timedays_range_str: - timemin, timemax = [float(timedays) for timedays in timedays_range_str.split("-")] + timemin, timemax = (float(timedays) for timedays in timedays_range_str.split("-")) else: timeavg = float(timedays_range_str) timestepmin = get_timestep_of_timedays(modelpath, timeavg) @@ -510,10 +535,10 @@ def get_escaped_arrivalrange(modelpath: Union[Path, str]) -> tuple[int, Optional # for 1D and 2D, the largest escape radius at tmin is the box side radius vmax_tmin = cornervmax if at.inputmodel.get_dfmodel_dimensions(dfmodel) == 3 else vmax - # earliest completely valid time is tmin plus maximum possible travel time + # earliest completely valid time is tmin plus maximum possible travel time from corner to origin validrange_start_days = at.get_timestep_times_float(modelpath, loc="start")[0] * (1 + vmax_tmin / 29979245800) - # find the last possible escape time and subtract the largest possible travel time + # find the last possible escape time and subtract the largest possible travel time (observer time correction) depdata = at.get_deposition(modelpath=modelpath) # use this file to find the last computed timestep nts_last = depdata.ts.max() if "ts" in depdata.columns else len(depdata) - 1 nts_last_tend = at.get_timestep_times_float(modelpath, loc="end")[nts_last] @@ -542,8 +567,8 @@ def get_model_name(path: Union[Path, str]) -> str: modelpath = abspath if os.path.isdir(abspath) else os.path.dirname(abspath) try: - plotlabelfile = os.path.join(modelpath, "plotlabel.txt") - return open(plotlabelfile, mode="r").readline().strip() + plotlabelfile = Path(modelpath, "plotlabel.txt") + return open(plotlabelfile).readline().strip() except FileNotFoundError: return os.path.basename(modelpath) @@ -560,9 +585,10 @@ def get_z_a_nucname(nucname: str) -> tuple[int, int]: @lru_cache(maxsize=1) def get_elsymbolslist() -> list[str]: - elsymbols = ["n"] + list( - pd.read_csv(at.get_config()["path_datadir"] / "elements.csv", usecols=["symbol"])["symbol"].values - ) + elsymbols = [ + "n", + *list(pd.read_csv(at.get_config()["path_datadir"] / "elements.csv", usecols=["symbol"])["symbol"].to_numpy()), + ] return elsymbols @@ -591,29 +617,32 @@ def get_elsymbol(atomic_number: int) -> str: @lru_cache(maxsize=16) -def get_ionstring(atomic_number: int, ionstage: int, spectral: bool = True, nospace: bool = False) -> str: +def get_ionstring( + atomic_number: int, ionstage: Union[int, Literal["ALL"], None], spectral: bool = True, nospace: bool = False +) -> str: if ionstage == "ALL" or ionstage is None: return f"{get_elsymbol(atomic_number)}" - elif spectral: + + if spectral: return f"{get_elsymbol(atomic_number)}{' ' if not nospace else ''}{roman_numerals[ionstage]}" + + # ion notion e.g. Co+, Fe2+ + if ionstage > 2: + strcharge = r"$^{" + str(ionstage - 1) + r"{+}}$" + elif ionstage == 2: + strcharge = r"$^{+}$" else: - # ion notion e.g. Co+, Fe2+ - if ionstage > 2: - strcharge = r"$^{" + str(ionstage - 1) + r"{+}}$" - elif ionstage == 2: - strcharge = r"$^{+}$" - else: - strcharge = "" - return f"{get_elsymbol(atomic_number)}{strcharge}" + strcharge = "" + return f"{get_elsymbol(atomic_number)}{strcharge}" # based on code from https://gist.github.com/kgaughan/2491663/b35e9a117b02a3567c8107940ac9b2023ba34ced -def parse_range(rng: str, dictvars: dict[str, int] = {}) -> Iterable[Any]: +def parse_range(rng: str, dictvars: dict[str, int]) -> Iterable[Any]: """Parse a string with an integer range and return a list of numbers, replacing special variables in dictvars.""" strparts = rng.split("-") if len(strparts) not in [1, 2]: - raise ValueError("Bad range: '%s'" % (rng,)) + raise ValueError(f"Bad range: '{rng}'") parts = [int(i) if i not in dictvars else dictvars[i] for i in strparts] start: int = parts[0] @@ -625,7 +654,7 @@ def parse_range(rng: str, dictvars: dict[str, int] = {}) -> Iterable[Any]: return range(start, end + 1) -def parse_range_list(rngs: Union[str, list], dictvars: dict = {}) -> list[Any]: +def parse_range_list(rngs: Union[str, list], dictvars: Optional[dict] = None) -> list[Any]: """Parse a string with comma-separated ranges or a list of range strings. Return a sorted list of integers in any of the ranges. @@ -635,19 +664,16 @@ def parse_range_list(rngs: Union[str, list], dictvars: dict = {}) -> list[Any]: elif not hasattr(rngs, "split"): return [rngs] - return sorted(set(chain.from_iterable([parse_range(rng, dictvars) for rng in rngs.split(",")]))) + return sorted(set(chain.from_iterable([parse_range(rng, dictvars or {}) for rng in rngs.split(",")]))) def makelist(x: Union[None, list, Sequence, str, Path]) -> list[Any]: """If x is not a list (or is a string), make a list containing x.""" if x is None: return [] - elif isinstance(x, (str, Path)): - return [ - x, - ] - else: - return list(x) + if isinstance(x, (str, Path)): + return [x] + return list(x) def trim_or_pad(requiredlength: int, *listoflistin: list[list[Any]]) -> Iterator[list[Any]]: @@ -671,62 +697,59 @@ def flatten_list(listin: list) -> list: return listout -def zopen(filename: Union[Path, str], mode: str): # type: ignore - """Open filename, filename.gz or filename.x""" - filenamexz = str(filename) if str(filename).endswith(".xz") else str(filename) + ".xz" - filenamegz = str(filename) if str(filename).endswith(".gz") else str(filename) + ".gz" - if os.path.exists(filename) and not str(filename).endswith(".gz") and not str(filename).endswith(".xz"): - return open(filename, mode) - elif os.path.exists(filenamegz) or str(filename).endswith(".gz"): - return gzip.open(filenamegz, mode) - elif os.path.exists(filenamexz) or str(filename).endswith(".xz"): - return lzma.open(filenamexz, mode) - else: - # will raise file not found - return open(filename, mode) +def zopen(filename: Union[Path, str], mode: str = "rt", encoding: Optional[str] = None) -> Any: + """Open filename, filename.gz or filename.xz""" + + ext_fopen = [(".lz4", lz4.frame.open), (".zst", pyzstd.open), (".gz", gzip.open), (".xz", xz.open)] + + for ext, fopen in ext_fopen: + file_ext = str(filename) if str(filename).endswith(ext) else str(filename) + ext + if Path(file_ext).exists(): + return fopen(file_ext, mode=mode, encoding=encoding) + + # open() can raise file not found if this file doesn't exist + return open(filename, mode=mode, encoding=encoding) def firstexisting( - filelist: Sequence[Union[str, Path]], - path: Union[Path, str] = Path("."), + filelist: Union[Sequence[Union[str, Path]], str, Path], + folder: Union[Path, str] = Path("."), tryzipped: bool = True, - noexcept: bool = False, ) -> Path: - """Return the first existing file in file list. If none exists, raise exception or if noexcept=True, return the last one""" - if isinstance(filelist, str) or isinstance(filelist, Path): + """Return the first existing file in file list. If none exist, raise exception.""" + if isinstance(filelist, (str, Path)): filelist = [filelist] fullpaths = [] for filename in filelist: - if tryzipped: - filenamexz = str(filename) if str(filename).endswith(".xz") else str(filename) + ".xz" - if filenamexz not in filelist: - fullpaths.append(Path(path) / filenamexz) + fullpaths.append(Path(folder) / filename) - filenamegz = str(filename) if str(filename).endswith(".gz") else str(filename) + ".gz" - if filenamegz not in filelist: - fullpaths.append(Path(path) / filenamegz) - - fullpaths.append(Path(path) / filename) + if tryzipped: + for ext in [".lz4", ".zst", ".gz", ".xz"]: + filenameext = str(filename) if str(filename).endswith(ext) else str(filename) + ext + if filenameext not in filelist: + fullpaths.append(Path(folder) / filenameext) for fullpath in fullpaths: if fullpath.exists(): return fullpath - if noexcept: - return fullpaths[-1] - - raise FileNotFoundError(f'None of these files exist in {path}: {", ".join([str(x) for x in fullpaths])}') + raise FileNotFoundError(f'None of these files exist in {folder}: {", ".join([str(x) for x in fullpaths])}') -def anyexist(filelist: Sequence[Union[str, Path]], path: Union[Path, str] = Path(".")) -> bool: - """Return the first existing file in file list.""" +def anyexist( + filelist: Sequence[Union[str, Path]], + folder: Union[Path, str] = Path("."), + tryzipped: bool = True, +) -> bool: + """Return true if any files in file list exist.""" - for fullpath in [Path(path) / filename for filename in filelist]: - if fullpath.exists(): - return True + try: + firstexisting(filelist=filelist, folder=folder, tryzipped=tryzipped) + except FileNotFoundError: + return False - return False + return True def stripallsuffixes(f: Path) -> Path: @@ -765,7 +788,7 @@ def add_derived_metadata(metadata: dict[str, Any]) -> dict[str, Any]: import yaml - filepath = Path(str(filepath).replace(".xz", "")) + filepath = Path(str(filepath).replace(".xz", "").replace(".gz", "").replace(".lz4", "").replace(".zst", "")) # check if the reference file (e.g. spectrum.txt) has an metadata file (spectrum.txt.meta.yml) individualmetafile = filepath.with_suffix(filepath.suffix + ".meta.yml") @@ -808,7 +831,7 @@ def movavgfilterfunc(ylist: Union[list[float], np.ndarray]) -> np.ndarray: if dictargs.get("filtersavgol", False): import scipy.signal - window_length, poly_order = [int(x) for x in args.filtersavgol] + window_length, poly_order = (int(x) for x in args.filtersavgol) def savgolfilterfunc(ylist: Union[list[float], np.ndarray]) -> np.ndarray: return scipy.signal.savgol_filter(ylist, window_length, poly_order, mode=mode) @@ -816,7 +839,7 @@ def savgolfilterfunc(ylist: Union[list[float], np.ndarray]) -> np.ndarray: assert filterfunc is None filterfunc = savgolfilterfunc - print("Applying Savitzky–Golay filter") + print("Applying Savitzky-Golay filter") return filterfunc @@ -827,7 +850,7 @@ def join_pdf_files(pdf_list: list[str], modelpath_list: list[Path]) -> None: merger = PdfFileMerger() for pdf, modelpath in zip(pdf_list, modelpath_list): - fullpath = firstexisting([pdf], path=modelpath) + fullpath = firstexisting([pdf], folder=modelpath) merger.append(open(fullpath, "rb")) os.remove(fullpath) @@ -842,17 +865,14 @@ def join_pdf_files(pdf_list: list[str], modelpath_list: list[Path]) -> None: def get_bflist(modelpath: Union[Path, str]) -> dict[int, tuple[int, int, int, int]]: compositiondata = get_composition_data(modelpath) bflist = {} - bflistpath = firstexisting(["bflist.out", "bflist.dat"], path=modelpath, tryzipped=True) - with zopen(bflistpath, "rt") as filein: + bflistpath = firstexisting(["bflist.out", "bflist.dat"], folder=modelpath, tryzipped=True) + with zopen(bflistpath) as filein: bflistcount = int(filein.readline()) - for k in range(bflistcount): + for _ in range(bflistcount): rowints = [int(x) for x in filein.readline().split()] i, elementindex, ionindex, level = rowints[:4] - if len(rowints) > 4: - upperionlevel = rowints[4] - else: - upperionlevel = -1 + upperionlevel = rowints[4] if len(rowints) > 4 else -1 atomic_number = compositiondata.Z[elementindex] ion_stage = ionindex + compositiondata.lowermost_ionstage[elementindex] bflist[i] = (atomic_number, ion_stage, level, upperionlevel) @@ -868,28 +888,28 @@ def read_linestatfile( ) -> tuple[int, list[float], list[int], list[int], list[int], list[int]]: """Load linestat.out containing transitions wavelength, element, ion, upper and lower levels.""" - with zopen(filepath, "rt") as linestatfile: - lambda_angstroms = [float(wl) * 1e8 for wl in linestatfile.readline().split()] - nlines = len(lambda_angstroms) + print(f"Loading {filepath}") + data = np.loadtxt(filepath) + lambda_angstroms = data[0] * 1e8 + nlines = len(lambda_angstroms) + + atomic_numbers = data[1].astype(int) + assert len(atomic_numbers) == nlines + + ion_stages = data[2].astype(int) + assert len(ion_stages) == nlines - atomic_numbers = [int(z) for z in linestatfile.readline().split()] - assert len(atomic_numbers) == nlines - ion_stages = [int(ion_stage) for ion_stage in linestatfile.readline().split()] - assert len(ion_stages) == nlines + # the file adds one to the levelindex, i.e. lowest level is 1 + upper_levels = data[3].astype(int) + assert len(upper_levels) == nlines - # the file adds one to the levelindex, i.e. lowest level is 1 - upper_levels = [int(levelplusone) - 1 for levelplusone in linestatfile.readline().split()] - assert len(upper_levels) == nlines - lower_levels = [int(levelplusone) - 1 for levelplusone in linestatfile.readline().split()] - assert len(lower_levels) == nlines + lower_levels = data[4].astype(int) + assert len(lower_levels) == nlines return nlines, lambda_angstroms, atomic_numbers, ion_stages, upper_levels, lower_levels -@lru_cache(maxsize=8) -def get_linelist_dict( - modelpath: Union[Path, str], returntype: Literal["dict", "dataframe"] = "dict" -) -> dict[int, linetuple]: +def get_linelist_dict(modelpath: Union[Path, str]) -> dict[int, linetuple]: nlines, lambda_angstroms, atomic_numbers, ion_stages, upper_levels, lower_levels = read_linestatfile( Path(modelpath, "linestat.out") ) @@ -910,8 +930,6 @@ def get_linelist_dataframe( Path(modelpath, "linestat.out") ) - # considering our standard lineline is about 1.5 million lines, - # using a dataframe make the lookup process very slow dflinelist = pd.DataFrame( { "lambda_angstroms": lambda_angstroms, @@ -919,7 +937,14 @@ def get_linelist_dataframe( "ionstage": ion_stages, "upperlevelindex": upper_levels, "lowerlevelindex": lower_levels, - } + }, + dtype={ + "lambda_angstroms": float, + "atomic_number": int, + "ionstage": int, + "upperlevelindex": int, + "lowerlevelindex": int, + }, ) dflinelist.index.name = "linelistindex" @@ -930,9 +955,11 @@ def get_linelist_dataframe( def get_npts_model(modelpath: Path) -> int: """Return the number of cell in the model.txt.""" modelfilepath = ( - Path(modelpath) if Path(modelpath).is_file() else at.firstexisting("model.txt", path=modelpath, tryzipped=True) + Path(modelpath) + if Path(modelpath).is_file() + else at.firstexisting("model.txt", folder=modelpath, tryzipped=True) ) - with zopen(modelfilepath, "rt") as modelfile: + with zopen(modelfilepath) as modelfile: npts_model = int(readnoncommentline(modelfile)) return npts_model @@ -946,8 +973,8 @@ def get_nprocs(modelpath: Path) -> int: @lru_cache(maxsize=8) def get_inputparams(modelpath: Path) -> dict[str, Any]: """Return parameters specified in input.txt.""" - from astropy import units as u from astropy import constants as const + from astropy import units as u params: dict[str, Any] = {} with Path(modelpath, "input.txt").open("r") as inputfile: @@ -957,21 +984,21 @@ def get_inputparams(modelpath: Path) -> dict[str, Any]: params["ntstep"] = int(readnoncommentline(inputfile).split("#")[0]) # number of start and end time step - params["itstep"], params["ftstep"] = [int(x) for x in readnoncommentline(inputfile).split("#")[0].split()] + params["itstep"], params["ftstep"] = (int(x) for x in readnoncommentline(inputfile).split("#")[0].split()) - params["tmin"], params["tmax"] = [float(x) for x in readnoncommentline(inputfile).split("#")[0].split()] + params["tmin"], params["tmax"] = (float(x) for x in readnoncommentline(inputfile).split("#")[0].split()) - params["nusyn_min"], params["nusyn_max"] = [ + params["nusyn_min"], params["nusyn_max"] = ( (float(x) * u.MeV / const.h).to("Hz") for x in readnoncommentline(inputfile).split("#")[0].split() - ] + ) # number of times for synthesis params["nsyn_time"] = int(readnoncommentline(inputfile).split("#")[0]) # start and end times for synthesis - params["nsyn_time_start"], params["nsyn_time_end"] = [ + params["nsyn_time_start"], params["nsyn_time_end"] = ( float(x) for x in readnoncommentline(inputfile).split("#")[0].split() - ] + ) params["n_dimensions"] = int(readnoncommentline(inputfile).split("#")[0]) @@ -985,7 +1012,7 @@ def get_runfolder_timesteps(folderpath: Union[Path, str]) -> tuple[int, ...]: """Get the set of timesteps covered by the output files in an ARTIS run folder.""" folder_timesteps = set() try: - with zopen(Path(folderpath, "estimators_0000.out"), "rt") as estfile: + with zopen(Path(folderpath, "estimators_0000.out")) as estfile: restart_timestep = -1 for line in estfile: if line.startswith("timestep "): @@ -1010,14 +1037,14 @@ def get_runfolders( """Get a list of folders containing ARTIS output files from a modelpath, optionally with a timestep restriction. The folder list may include non-ARTIS folders if a timestep is not specified.""" - folderlist_all = tuple(sorted([child for child in Path(modelpath).iterdir() if child.is_dir()]) + [Path(modelpath)]) + folderlist_all = (*sorted([child for child in Path(modelpath).iterdir() if child.is_dir()]), Path(modelpath)) folder_list_matching = [] if (timestep is not None and timestep > -1) or (timesteps is not None and len(timesteps) > 0): for folderpath in folderlist_all: folder_timesteps = get_runfolder_timesteps(folderpath) if timesteps is None and timestep is not None and timestep in folder_timesteps: return (folderpath,) - elif timesteps is not None and any([ts in folder_timesteps for ts in timesteps]): + if timesteps is not None and any(ts in folder_timesteps for ts in timesteps): folder_list_matching.append(folderpath) return tuple(folder_list_matching) @@ -1039,29 +1066,29 @@ def get_mpiranklist( - only_ranks_withgridcells: set True to skip ranks that only update packets (i.e. that don't update any grid cells/output estimators) """ + if modelgridindex is None or modelgridindex == []: if only_ranks_withgridcells: return range(min(get_nprocs(modelpath), get_npts_model(modelpath))) return range(get_nprocs(modelpath)) - else: - if isinstance(modelgridindex, Iterable): - mpiranklist = set() - for mgi in modelgridindex: - if mgi < 0: - if only_ranks_withgridcells: - return range(min(get_nprocs(modelpath), get_npts_model(modelpath))) - return range(get_nprocs(modelpath)) - else: - mpiranklist.add(get_mpirankofcell(mgi, modelpath=modelpath)) - - return sorted(list(mpiranklist)) - else: - # in case modelgridindex is a single number rather than an iterable - if modelgridindex < 0: - return range(min(get_nprocs(modelpath), get_npts_model(modelpath))) - else: - return [get_mpirankofcell(modelgridindex, modelpath=modelpath)] + if isinstance(modelgridindex, Iterable): + mpiranklist = set() + for mgi in modelgridindex: + if mgi < 0: + if only_ranks_withgridcells: + return range(min(get_nprocs(modelpath), get_npts_model(modelpath))) + return range(get_nprocs(modelpath)) + + mpiranklist.add(get_mpirankofcell(mgi, modelpath=modelpath)) + + return sorted(mpiranklist) + + # in case modelgridindex is a single number rather than an iterable + if modelgridindex < 0: + return range(min(get_nprocs(modelpath), get_npts_model(modelpath))) + + return [get_mpirankofcell(modelgridindex, modelpath=modelpath)] def get_cellsofmpirank(mpirank: int, modelpath: Union[Path, str]) -> Iterable[int]: @@ -1089,7 +1116,7 @@ def get_dfrankassignments(modelpath: Union[Path, str]) -> Optional[pd.DataFrame] filerankassignments = Path(modelpath, "modelgridrankassignments.out") if filerankassignments.is_file(): df = pd.read_csv(filerankassignments, delim_whitespace=True) - df.rename(columns={df.columns[0]: df.columns[0].lstrip("#")}, inplace=True) + df = df.rename(columns={df.columns[0]: df.columns[0].lstrip("#")}) return df return None @@ -1116,133 +1143,17 @@ def get_mpirankofcell(modelgridindex: int, modelpath: Union[Path, str]) -> int: nblock = npts_model // nprocs n_leftover = npts_model % nprocs - if modelgridindex <= n_leftover * (nblock + 1): - mpirank = modelgridindex // (nblock + 1) - else: - mpirank = n_leftover + (modelgridindex - n_leftover * (nblock + 1)) // nblock + mpirank = ( + modelgridindex // (nblock + 1) + if modelgridindex <= n_leftover * (nblock + 1) + else n_leftover + (modelgridindex - n_leftover * (nblock + 1)) // nblock + ) assert modelgridindex in get_cellsofmpirank(mpirank, modelpath) return mpirank -def get_artis_constants( - modelpath: Optional[Path] = None, srcpath: Optional[Path] = None, printdefs: bool = False -) -> dict[str, Any]: - # get artis options specified as preprocessor macro definitions in artisoptions.h and other header files - if not srcpath: - assert modelpath is not None - srcpath = Path(modelpath, "artis") - if not modelpath: - raise ValueError("Either modelpath or srcpath must be specified in call to get_defines()") - - cfiles = [ - # Path(srcpath, 'constants.h'), - # Path(srcpath, 'decay.h'), - Path(srcpath, "artisoptions.h"), - # Path(srcpath, 'sn3d.h'), - ] - definedict = { - "true": True, - "false": False, - } - for filepath in cfiles: - definedict.update(parse_cdefines(srcfilepath=filepath)) - - # evaluate booleans, numbers, and references to other constants - for k, strvalue in definedict.copy().items(): - try: - # definedict[k] = eval(strvalue, definedict) - # print(f"{k} = '{strvalue}' = {definedict[k]}") - pass - except SyntaxError: - pass - # print(f"{k} = '{strvalue}' = (COULD NOT EVALUATE)") - except TypeError: - pass - # print(f"{k} = '{strvalue}' = (COULD NOT EVALUATE)") - - # if printdefs: - # for k in definedict: - # print(f"{k} = '{definedict[k]}'") - - return definedict - - -def parse_cdefines(srcfilepath: Path, printdefs: bool = False) -> dict[str, Any]: - # adapted from h2py.py in Python source - import re - - # p_define = re.compile('^[\t ]*#[\t ]*define[\t ]+([a-zA-Z0-9_]+)[\t ]+') - p_define = re.compile(r"^[\t ]*#[\t ]*define[\t ]+([a-zA-Z0-9_]+)+") - - p_const = re.compile(r"(?:\w+\s+)([a-zA-Z_=][a-zA-Z0-9_=]*)*(? str: - # replace ignored patterns by spaces - for p in ignores: - body = p.sub(" ", body) - # replace char literals by ord(...) - body = p_char.sub("ord(\\0)", body) - # Compute negative hexadecimal constants - start = 0 - UMAX = 2 * (sys.maxsize + 1) - while 1: - m = p_hex.search(body, start) - if not m: - break - s, e = m.span() - val = int(body[slice(*m.span(1))], 16) - if val > sys.maxsize: - val -= UMAX - body = body[:s] + "(" + str(val) + ")" + body[e:] - start = s + 1 - return body - - definedict = {} - lineno = 0 - with open(srcfilepath, "r") as optfile: - while 1: - line = optfile.readline() - if not line: - break - lineno = lineno + 1 - match = p_define.match(line) - if match: - # gobble up continuation lines - while line[-2:] == "\\\n": - nextline = optfile.readline() - if not nextline: - break - lineno = lineno + 1 - line = line + nextline - name = match.group(1) - body = line[match.end() :] - body = pytify(body) - definedict[name] = body.strip() - match = p_const.match(line) - if match: - print("CONST", tuple(p_const.findall(line))) - # if '=' in line and ';' in line: - # tokens = line.replace('==', 'IGNORE').replace('=', ' = ').split() - # varname = tokens.indexof('=')[-1] - - if printdefs: - for k in definedict: - print(f"{k} = '{definedict[k]}'") - - return definedict - - def get_viewingdirectionbincount() -> int: return 100 @@ -1255,23 +1166,23 @@ def get_viewingdirection_costhetabincount() -> int: return 10 -def get_viewinganglebin_definitions() -> tuple[list[str], list[str]]: +def get_costhetabin_phibin_labels() -> tuple[list[str], list[str]]: # todo: replace with general code for any bin count: # ncosthetabins = at.get_viewingdirection_costhetabincount() # costhetabins_lower = np.arange(-1., 1., 2. / ncosthetabins) # costhetabins_upper = costhetabins_lower + 2. / ncosthetabins costheta_viewing_angle_bins = [ - "-1.0 ≤ cos(θ) < -0.8", - "-0.8 ≤ cos(θ) < -0.6", - "-0.6 ≤ cos(θ) < -0.4", - "-0.4 ≤ cos(θ) < -0.2", - "-0.2 ≤ cos(θ) < 0.0", - " 0.0 ≤ cos(θ) < 0.2", - " 0.2 ≤ cos(θ) < 0.4", - " 0.4 ≤ cos(θ) < 0.6", - " 0.6 ≤ cos(θ) < 0.8", - " 0.8 ≤ cos(θ) < 1.0", + "-1.0 ≤ cos θ < -0.8", + "-0.8 ≤ cos θ < -0.6", + "-0.6 ≤ cos θ < -0.4", + "-0.4 ≤ cos θ < -0.2", + "-0.2 ≤ cos θ < 0.0", + " 0.0 ≤ cos θ < 0.2", + " 0.2 ≤ cos θ < 0.4", + " 0.4 ≤ cos θ < 0.6", + " 0.6 ≤ cos θ < 0.8", + " 0.8 ≤ cos θ < 1.0", ] assert len(costheta_viewing_angle_bins) == get_viewingdirection_costhetabincount() @@ -1302,3 +1213,55 @@ def get_viewinganglebin_definitions() -> tuple[list[str], list[str]]: # '6π/5 < ϕ ≤ 7π/5', '7π/5 < ϕ ≤ 8π/5', # '8π/5 < ϕ ≤ 9π/5', '9π/5 < ϕ < 2π'] return costheta_viewing_angle_bins, phi_viewing_angle_bins + + +def get_vspec_dir_labels(modelpath: Union[str, Path], viewinganglelabelunits: str = "rad") -> dict[int, str]: + vpkt_config = at.get_vpkt_config(modelpath) + dirlabels = {} + for dirindex in range(vpkt_config["nobsdirections"]): + phi_angle = round(vpkt_config["phi"][dirindex]) + if viewinganglelabelunits == "deg": + theta_angle = round(math.degrees(math.acos(vpkt_config["cos_theta"][dirindex]))) + dirlabels[dirindex] = rf"v$\theta$ = {theta_angle}$^\circ$, $\phi$ = {phi_angle}$^\circ$" + elif viewinganglelabelunits == "rad": + dirlabels[dirindex] = rf"cos $\theta$ = {vpkt_config['cos_theta'][dirindex]}, $\phi$ = {phi_angle}$^\circ$" + return dirlabels + + +def get_dirbin_labels( + dirbins: Union[np.ndarray[Any, np.dtype[Any]], Sequence[int]], + modelpath: Union[Path, str, None] = None, + average_over_phi: bool = False, + average_over_theta: bool = False, +) -> dict[int, str]: + if modelpath: + modelpath = Path(modelpath) + MABINS = at.get_viewingdirectionbincount() + if len(list(Path(modelpath).glob("*_res_00.out*"))) > 0: # if the first direction bin file exists + assert len(list(Path(modelpath).glob(f"*_res_{MABINS-1:02d}.out*"))) > 0 # check last bin exists + assert len(list(Path(modelpath).glob(f"*_res_{MABINS:02d}.out*"))) == 0 # check one beyond does not exist + + strlist_costheta_bins, strlist_phi_bins = at.get_costhetabin_phibin_labels() + + nphibins = at.get_viewingdirection_phibincount() + + angle_definitions: dict[int, str] = {} + for dirbin in dirbins: + if dirbin == -1: + angle_definitions[dirbin] = "" + continue + + costheta_index = dirbin // nphibins + phi_index = dirbin % nphibins + + if average_over_phi: + angle_definitions[dirbin] = f"{strlist_costheta_bins[costheta_index]}" + assert phi_index == 0 + assert not average_over_theta + elif average_over_theta: + angle_definitions[dirbin] = f"{strlist_phi_bins[phi_index]}" + assert costheta_index == 0 + else: + angle_definitions[dirbin] = f"{strlist_costheta_bins[costheta_index]}, {strlist_phi_bins[phi_index]}" + + return angle_definitions diff --git a/artistools/nltepops/__init__.py b/artistools/nltepops/__init__.py index 486120a58..75213c587 100644 --- a/artistools/nltepops/__init__.py +++ b/artistools/nltepops/__init__.py @@ -1,10 +1,10 @@ -#!/usr/bin/env python3 -"""Artistools - light curve functions.""" -from artistools.nltepops.nltepops import add_lte_pops -from artistools.nltepops.nltepops import read_file -from artistools.nltepops.nltepops import read_file_filtered -from artistools.nltepops.nltepops import read_files -from artistools.nltepops.nltepops import texifyconfiguration -from artistools.nltepops.nltepops import texifyterm -from artistools.nltepops.plotnltepops import addargs -from artistools.nltepops.plotnltepops import main +"""Artistools - non-LTE population functions.""" +from .__main__ import main +from .nltepops import add_lte_pops +from .nltepops import read_file +from .nltepops import read_file_filtered +from .nltepops import read_files +from .nltepops import texifyconfiguration +from .nltepops import texifyterm +from .plotnltepops import addargs +from .plotnltepops import main as plot diff --git a/artistools/nltepops/__main__.py b/artistools/nltepops/__main__.py index 9e73789a9..408583ce1 100644 --- a/artistools/nltepops/__main__.py +++ b/artistools/nltepops/__main__.py @@ -1,6 +1,9 @@ -import artistools as at -import artistools.nltepops +from .plotnltepops import main as plot + + +def main() -> None: + plot() + if __name__ == "__main__": - # multiprocessing.freeze_support() - at.nltepops.main() + main() diff --git a/artistools/nltepops/nltepops.py b/artistools/nltepops/nltepops.py index e60dbc3a6..b45beb742 100644 --- a/artistools/nltepops/nltepops.py +++ b/artistools/nltepops/nltepops.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 """Artistools - NLTE population related functions.""" import math import multiprocessing @@ -6,24 +5,17 @@ from functools import lru_cache from functools import partial from pathlib import Path +from typing import Optional +from typing import Union import pandas as pd from astropy import constants as const import artistools as at -import artistools.misc +from artistools.configuration import get_config -# import os -# import sys -# from itertools import chain -# import matplotlib.pyplot as plt -# import matplotlib.ticker as ticker -# import numpy as np -# import matplotlib as mpl - - -def texifyterm(strterm): +def texifyterm(strterm: str) -> str: """Replace a term string with TeX notation equivalent.""" strtermtex = "" passed_term_Lchar = False @@ -51,7 +43,7 @@ def texifyterm(strterm): return strtermtex -def texifyconfiguration(levelname): +def texifyconfiguration(levelname: str) -> str: """Replace a level configuration with the formatted LaTeX equivalent.""" # the underscore gets confused with LaTeX subscript operator, so switch it to the hash symbol strout = "#".join(levelname.split("_")[:-1]) + "#" @@ -113,7 +105,7 @@ def add_lte_pops(modelpath, dfpop, columntemperature_tuples, noprint=False, maxl & (dfpop["level"] != -1) ) - def f_ltepop(x, T_exc, gsg, gse, ionlevels): + def f_ltepop(x, T_exc: float, gsg: float, gse: float, ionlevels) -> float: return ( ionlevels.iloc[int(x.level)].g / gsg @@ -153,11 +145,10 @@ def f_ltepop(x, T_exc, gsg, gse, ionlevels): return dfpop -@at.diskcache(savezipped=True) -def read_file(nltefilepath): +def read_file(nltefilepath: Union[str, Path]) -> pd.DataFrame: """Read NLTE populations from one file.""" - if not nltefilepath.is_file(): + if not Path(nltefilepath).is_file(): nltefilepathgz = Path(str(nltefilepath) + ".gz") nltefilepathxz = Path(str(nltefilepath) + ".xz") if nltefilepathxz.is_file(): @@ -183,16 +174,20 @@ def read_file_filtered(nltefilepath, strquery=None, dfqueryvars=None): dfpopfile = read_file(nltefilepath) if strquery and not dfpopfile.empty: - dfpopfile.query(strquery, local_dict=dfqueryvars, inplace=True) + dfpopfile = dfpopfile.query(strquery, local_dict=dfqueryvars) return dfpopfile @lru_cache(maxsize=2) -@at.diskcache(savezipped=True, funcversion="2020-07-03.1327", saveonly=False) -def read_files(modelpath, timestep=-1, modelgridindex=-1, dfquery=None, dfqueryvars={}): +def read_files( + modelpath, timestep=-1, modelgridindex=-1, dfquery=None, dfqueryvars: Optional[dict] = None +) -> pd.DataFrame: """Read in NLTE populations from a model for a particular timestep and grid cell.""" + if dfqueryvars is None: + dfqueryvars = {} + mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=modelgridindex) dfpop = pd.DataFrame() @@ -220,8 +215,8 @@ def read_files(modelpath, timestep=-1, modelgridindex=-1, dfquery=None, dfqueryv dfquery_full = f"({dfquery_full}) and " dfquery_full += f"({dfquery})" - if at.get_config()["num_processes"] > 1: - with multiprocessing.Pool(processes=at.get_config()["num_processes"]) as pool: + if get_config()["num_processes"] > 1: + with multiprocessing.Pool(processes=get_config()["num_processes"]) as pool: arr_dfnltepop = pool.map( partial(read_file_filtered, strquery=dfquery_full, dfqueryvars=dfqueryvars), nltefilepaths ) diff --git a/artistools/nltepops/plotnltepops.py b/artistools/nltepops/plotnltepops.py index 9785c8648..3532ce5db 100755 --- a/artistools/nltepops/plotnltepops.py +++ b/artistools/nltepops/plotnltepops.py @@ -4,28 +4,17 @@ import math import multiprocessing import os +import sys from pathlib import Path import matplotlib as mpl import matplotlib.pyplot as plt -import matplotlib.ticker as ticker import numpy as np import pandas as pd from astropy import constants as const +from matplotlib import ticker import artistools as at -import artistools.atomic -import artistools.estimators -import artistools.nltepops -import artistools.plottools - -# import re -# import sys -# from functools import lru_cache -# from functools import partial -# from itertools import chain -# import numpy as np - defaultoutputfile = "plotnlte_{elsymbol}_cell{cell:03d}_ts{timestep:02d}_{time_days:.0f}d.pdf" @@ -37,7 +26,7 @@ def annotate_emission_line(ax, y, upperlevel, lowerlevel, label): xycoords=("data", "axes fraction"), xytext=(upperlevel, y), textcoords=("data", "axes fraction"), - arrowprops=dict(facecolor="black", width=0.1, headwidth=6), + arrowprops={"facecolor": "black", "width": 0.1, "headwidth": 6}, ) ax.annotate( @@ -50,8 +39,8 @@ def annotate_emission_line(ax, y, upperlevel, lowerlevel, label): ) -def plot_reference_data(ax, atomic_number, ion_stage, estimators_celltimestep, dfpopthision, args, annotatelines): - nne, Te, TR, W = [estimators_celltimestep[s] for s in ["nne", "Te", "TR", "W"]] +def plot_reference_data(ax, atomic_number, ion_stage, estimators_celltimestep, dfpopthision, annotatelines): + nne, Te, TR, W = (estimators_celltimestep[s] for s in ["nne", "Te", "TR", "W"]) # comparison to Chianti file elsym = at.get_elsymbol(atomic_number) elsymlower = elsym.lower() @@ -76,10 +65,10 @@ def plot_reference_data(ax, atomic_number, ion_stage, estimators_celltimestep, d if math.isclose(file_nne, nne, rel_tol=0.01) and math.isclose(file_Te, Te, abs_tol=10): if file_W > 0: continue - bbstr = " with dilute blackbody" - color = "C2" - marker = "+" - else: + # bbstr = " with dilute blackbody" + # color = "C2" + # marker = "+" + else: # noqa: RET507 bbstr = "" color = "C1" marker = "^" @@ -96,7 +85,7 @@ def plot_reference_data(ax, atomic_number, ion_stage, estimators_celltimestep, d row = line.split() try: levelnum = levelnumofconfigterm[(row[1], row[2])] - if levelnum in dfpopthision["level"].values: + if levelnum in dfpopthision["level"].to_numpy(): levelnums.append(levelnum) if firstdep < 0: firstdep = float(row[0]) @@ -135,9 +124,9 @@ def get_floers_data(dfpopthision, atomic_number, ion_stage, modelpath, T_e, mode print(f"reading {floersfilename}") floers_levelpops = pd.read_csv(modelpath / floersfilename, comment="#", delim_whitespace=True) # floers_levelnums = floers_levelpops['index'].values - 1 - floers_levelpops.sort_values(by="energypercm", inplace=True) + floers_levelpops = floers_levelpops.sort_values(by="energypercm") floers_levelnums = list(range(len(floers_levelpops))) - floers_levelpop_values = floers_levelpops["frac_ionpop"].values * dfpopthision["n_NLTE"].sum() + floers_levelpop_values = floers_levelpops["frac_ionpop"].to_numpy() * dfpopthision["n_NLTE"].sum() floersmultizonefilename = None if modelpath.stem.startswith("w7_"): @@ -167,7 +156,7 @@ def get_floers_data(dfpopthision, atomic_number, ion_stage, modelpath, T_e, mode if abs(row["vel_outer"] - vel_outer) < 0.5: print(f" ARTIS cell vel_outter: {vel_outer}, Floersfile: {row['vel_outer']}") print(f" ARTIS cell Te: {T_e}, Floersfile: {row['Te']}") - floers_levelpops = row.values[4:] + floers_levelpops = row.to_numpy()[4:] if len(dfpopthision["level"]) < len(floers_levelpops): floers_levelpops = floers_levelpops[: len(dfpopthision["level"])] floers_levelnums = list(range(len(floers_levelpops))) @@ -209,7 +198,7 @@ def make_ionsubplot( dfpopthision = at.nltepops.add_lte_pops(modelpath, dfpopthision, lte_columns, noprint=False, maxlevel=args.maxlevel) if args.maxlevel >= 0: - dfpopthision.query("level <= @args.maxlevel", inplace=True) + dfpopthision = dfpopthision.query("level <= @args.maxlevel") ionpopulation = dfpopthision["n_NLTE"].sum() ionpopulation_fromest = estimators[(timestep, modelgridindex)]["populations"].get((atomic_number, ion_stage), 0.0) @@ -255,16 +244,17 @@ def make_ionsubplot( f"level population of {ionpopulation:.1f} (from estimator file ion pop = {ionpopulation_fromest})" ) - if args.departuremode: + lte_scalefactor = ( # scale to match the ground state populations - lte_scalefactor = float(dfpopthision["n_NLTE"].iloc[0] / dfpopthision["n_LTE_T_e"].iloc[0]) - else: - # scale to match the ion population - lte_scalefactor = float(ionpopulation / dfpopthision["n_LTE_T_e"].sum()) + float(dfpopthision["n_NLTE"].iloc[0] / dfpopthision["n_LTE_T_e"].iloc[0]) + if args.departuremode + # else scale to match the ion population + else float(ionpopulation / dfpopthision["n_LTE_T_e"].sum()) + ) - dfpopthision.eval("n_LTE_T_e_normed = n_LTE_T_e * @x", local_dict={"x": lte_scalefactor}, inplace=True) + dfpopthision = dfpopthision.eval("n_LTE_T_e_normed = n_LTE_T_e * @x", local_dict={"x": lte_scalefactor}) - dfpopthision.eval("departure_coeff = n_NLTE / n_LTE_T_e_normed", inplace=True) + dfpopthision = dfpopthision.eval("departure_coeff = n_NLTE / n_LTE_T_e_normed") pd.set_option("display.max_columns", 150) if len(dfpopthision) < 30: @@ -295,7 +285,7 @@ def make_ionsubplot( round((const.h * const.c).to("eV angstrom").value / trans.energy_trans) for _, trans in dftrans.iterrows() ] - dftrans.sort_values("emissionstrength", ascending=False, inplace=True) + dftrans = dftrans.sort_values("emissionstrength", ascending=False) print("\nTop radiative decays") print(dftrans[:10].to_string(index=False)) @@ -349,7 +339,7 @@ def make_ionsubplot( if not args.hide_lte_tr: lte_scalefactor = float(ionpopulation / dfpopthision["n_LTE_T_R"].sum()) - dfpopthision.eval("n_LTE_T_R_normed = n_LTE_T_R * @lte_scalefactor", inplace=True) + dfpopthision = dfpopthision.eval("n_LTE_T_R_normed = n_LTE_T_R * @lte_scalefactor") ax.plot( dfpopthision["level"], dfpopthision["n_LTE_T_R_normed"], @@ -385,7 +375,7 @@ def make_ionsubplot( if args.plotrefdata: plot_reference_data( - ax, atomic_number, ion_stage, estimators[(timestep, modelgridindex)], dfpopthision, args, annotatelines=True + ax, atomic_number, ion_stage, estimators[(timestep, modelgridindex)], dfpopthision, annotatelines=True ) @@ -425,10 +415,7 @@ def make_plot_populations_with_time_or_velocity(modelpaths, args): ax = ax.flatten() for plotnumber, timedays in enumerate(timedayslist): - if args.subplots: - axis = ax[plotnumber] - else: - axis = ax + axis = ax[plotnumber] if args.subplots else ax plot_populations_with_time_or_velocity( axis, modelpaths, timedays, ionstage, ionlevels, Z, levelconfignames, args ) @@ -440,8 +427,6 @@ def make_plot_populations_with_time_or_velocity(modelpaths, args): xlabel = r"Zone outer velocity [km s$^{-1}$]" ylabel = r"Level population [cm$^{-3}$]" - import artistools.plottools - at.plottools.set_axis_labels(fig, ax, xlabel, ylabel, labelfontsize, args) if args.subplots: for plotnumber, axis in enumerate(ax): @@ -472,11 +457,11 @@ def make_plot_populations_with_time_or_velocity(modelpaths, args): def plot_populations_with_time_or_velocity(ax, modelpaths, timedays, ionstage, ionlevels, Z, levelconfignames, args): if args.x == "time": - timesteps = [time for time in range(args.timestepmin, args.timestepmax)] + timesteps = list(range(args.timestepmin, args.timestepmax)) if not args.modelgridindex: print("Please specify modelgridindex") - quit() + sys.exit(1) modelgridindex_list = np.ones_like(timesteps) modelgridindex_list = modelgridindex_list * int(args.modelgridindex[0]) @@ -505,15 +490,15 @@ def plot_populations_with_time_or_velocity(ax, modelpaths, timedays, ionstage, i for ionlevel in ionlevels: populations[(timestep, ionlevel, mgi)] = timesteppops.loc[timesteppops["level"] == ionlevel][ "n_NLTE" - ].values[0] + ].to_numpy()[0] # populationsLTE[(timestep, ionlevel)] = (timesteppops.loc[timesteppops['level'] # == ionlevel]['n_LTE'].values[0]) for ionlevel in ionlevels: - plottimesteps = np.array([int(ts) for ts, level, mgi in populations.keys() if level == ionlevel]) + plottimesteps = np.array([int(ts) for ts, level, mgi in populations if level == ionlevel]) timedays = [float(at.get_timestep_time(modelpath, ts)) for ts in plottimesteps] plotpopulations = np.array( - [float(populations[ts, level, mgi]) for ts, level, mgi in populations.keys() if level == ionlevel] + [float(populations[ts, level, mgi]) for ts, level, mgi in populations if level == ionlevel] ) # plotpopulationsLTE = np.array([float(populationsLTE[ts, level]) for ts, level in populationsLTE.keys() # if level == ionlevel]) @@ -542,7 +527,7 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep, print(f"No NLTE population data for modelgrid cell {mgilist[0]} timestep {timestep}") return - dfpop.query("Z == @atomic_number", inplace=True) + dfpop = dfpop.query("Z == @atomic_number") # top_ion = 9999 max_ion_stage = dfpop.ion_stage.max() @@ -576,6 +561,7 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep, axes = [axes] prev_ion_stage = -1 + assert len(mgilist) > 0 for mgilistindex, modelgridindex in enumerate(mgilist): mgifirstaxindex = mgilistindex mgilastaxindex = mgilistindex + len(ion_stage_list) - 1 @@ -609,7 +595,7 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep, print(f"No NLTE population data for modelgrid cell {modelgridindex} timestep {timestep}") return - dfpop.query("Z == @atomic_number", inplace=True) + dfpop = dfpop.query("Z == @atomic_number") # top_ion = 9999 max_ion_stage = dfpop.ion_stage.max() @@ -634,7 +620,7 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep, subplot_title += f" timestep {timestep:d}" else: subplot_title += f" {time_days:.0f}d" - subplot_title += f" (Te={T_e:.0f} K, nne={nne:.1e} " + r"cm$^{-3}$, T$_R$=" + f"{T_R:.0f} K, W={W:.1e})" + subplot_title += rf" (Te={T_e:.0f} K, nne={nne:.1e} cm$^{-3}$, T$_R$={T_R:.0f} K, W={W:.1e})" if not args.notitle: axes[mgifirstaxindex].set_title(subplot_title, fontsize=10) @@ -681,7 +667,7 @@ def make_plot(modelpath, atomic_number, ionstages_displayed, mgilist, timestep, axes[-1].set_xlabel(r"Level index") outputfilename = str(args.outputfile).format( - elsymbol=at.get_elsymbol(atomic_number), cell=modelgridindex, timestep=timestep, time_days=time_days + elsymbol=at.get_elsymbol(atomic_number), cell=mgilist[0], timestep=timestep, time_days=time_days ) fig.savefig(str(outputfilename), format="pdf") print(f"Saved {outputfilename}") @@ -759,6 +745,8 @@ def main(args=None, argsraw=None, **kwargs): parser.set_defaults(**kwargs) args = parser.parse_args(argsraw) + at.set_mpl_style() + if args.x in ["time", "velocity"]: # if len(args.modelpath) == 1: # modelpath = args.modelpath @@ -767,13 +755,13 @@ def main(args=None, argsraw=None, **kwargs): # if not args.timedays: # print("Please specify time range with -timedays") - # quit() + # sys.exit(1) if not args.ionstages: print("Please specify ionstage") - quit() + sys.exit(1) if not args.levels: print("Please specify levels") - quit() + sys.exit(1) else: modelpath = args.modelpath @@ -801,7 +789,7 @@ def main(args=None, argsraw=None, **kwargs): if isinstance(args.elements, str): args.elements = [args.elements] - if isinstance(args.velocity, float) or isinstance(args.velocity, int): + if isinstance(args.velocity, (float, int)): args.velocity = [args.velocity] mgilist = [] diff --git a/artistools/nonthermal/__init__.py b/artistools/nonthermal/__init__.py index 6cc05f200..1d7662891 100644 --- a/artistools/nonthermal/__init__.py +++ b/artistools/nonthermal/__init__.py @@ -1,44 +1,44 @@ -#!/usr/bin/env python3 """Artistools - spectra related functions.""" -from artistools.nonthermal._nonthermal_core import analyse_ntspectrum -from artistools.nonthermal._nonthermal_core import ar_xs -from artistools.nonthermal._nonthermal_core import calculate_frac_heating -from artistools.nonthermal._nonthermal_core import calculate_Latom_excitation -from artistools.nonthermal._nonthermal_core import calculate_Latom_ionisation -from artistools.nonthermal._nonthermal_core import calculate_N_e -from artistools.nonthermal._nonthermal_core import calculate_nt_frac_excitation -from artistools.nonthermal._nonthermal_core import differentialsfmatrix_add_ionization_shell -from artistools.nonthermal._nonthermal_core import e_s_test -from artistools.nonthermal._nonthermal_core import get_arxs_array_ion -from artistools.nonthermal._nonthermal_core import get_arxs_array_shell -from artistools.nonthermal._nonthermal_core import get_electronoccupancy -from artistools.nonthermal._nonthermal_core import get_energyindex_gteq -from artistools.nonthermal._nonthermal_core import get_energyindex_lteq -from artistools.nonthermal._nonthermal_core import get_epsilon_avg -from artistools.nonthermal._nonthermal_core import get_fij_ln_en_ionisation -from artistools.nonthermal._nonthermal_core import get_J -from artistools.nonthermal._nonthermal_core import get_Latom_axelrod -from artistools.nonthermal._nonthermal_core import get_Lelec_axelrod -from artistools.nonthermal._nonthermal_core import get_lotz_xs_ionisation -from artistools.nonthermal._nonthermal_core import get_mean_binding_energy -from artistools.nonthermal._nonthermal_core import get_mean_binding_energy_alt -from artistools.nonthermal._nonthermal_core import get_nne -from artistools.nonthermal._nonthermal_core import get_nne_nt -from artistools.nonthermal._nonthermal_core import get_nnetot -from artistools.nonthermal._nonthermal_core import get_nntot -from artistools.nonthermal._nonthermal_core import get_xs_excitation -from artistools.nonthermal._nonthermal_core import get_xs_excitation_vector -from artistools.nonthermal._nonthermal_core import get_Zbar -from artistools.nonthermal._nonthermal_core import get_Zboundbar -from artistools.nonthermal._nonthermal_core import lossfunction -from artistools.nonthermal._nonthermal_core import lossfunction_axelrod -from artistools.nonthermal._nonthermal_core import namedtuple -from artistools.nonthermal._nonthermal_core import Psecondary -from artistools.nonthermal._nonthermal_core import read_binding_energies -from artistools.nonthermal._nonthermal_core import read_colliondata -from artistools.nonthermal._nonthermal_core import sfmatrix_add_excitation -from artistools.nonthermal._nonthermal_core import sfmatrix_add_ionization_shell -from artistools.nonthermal._nonthermal_core import solve_spencerfano_differentialform -from artistools.nonthermal._nonthermal_core import workfunction_tests -from artistools.nonthermal.plotnonthermal import addargs -from artistools.nonthermal.plotnonthermal import main +from .__main__ import main +from ._nonthermal_core import analyse_ntspectrum +from ._nonthermal_core import ar_xs +from ._nonthermal_core import calculate_frac_heating +from ._nonthermal_core import calculate_Latom_excitation +from ._nonthermal_core import calculate_Latom_ionisation +from ._nonthermal_core import calculate_N_e +from ._nonthermal_core import calculate_nt_frac_excitation +from ._nonthermal_core import differentialsfmatrix_add_ionization_shell +from ._nonthermal_core import e_s_test +from ._nonthermal_core import get_arxs_array_ion +from ._nonthermal_core import get_arxs_array_shell +from ._nonthermal_core import get_electronoccupancy +from ._nonthermal_core import get_energyindex_gteq +from ._nonthermal_core import get_energyindex_lteq +from ._nonthermal_core import get_epsilon_avg +from ._nonthermal_core import get_fij_ln_en_ionisation +from ._nonthermal_core import get_J +from ._nonthermal_core import get_Latom_axelrod +from ._nonthermal_core import get_Lelec_axelrod +from ._nonthermal_core import get_lotz_xs_ionisation +from ._nonthermal_core import get_mean_binding_energy +from ._nonthermal_core import get_mean_binding_energy_alt +from ._nonthermal_core import get_nne +from ._nonthermal_core import get_nne_nt +from ._nonthermal_core import get_nnetot +from ._nonthermal_core import get_nntot +from ._nonthermal_core import get_xs_excitation +from ._nonthermal_core import get_xs_excitation_vector +from ._nonthermal_core import get_Zbar +from ._nonthermal_core import get_Zboundbar +from ._nonthermal_core import lossfunction +from ._nonthermal_core import lossfunction_axelrod +from ._nonthermal_core import namedtuple +from ._nonthermal_core import Psecondary +from ._nonthermal_core import read_binding_energies +from ._nonthermal_core import read_colliondata +from ._nonthermal_core import sfmatrix_add_excitation +from ._nonthermal_core import sfmatrix_add_ionization_shell +from ._nonthermal_core import solve_spencerfano_differentialform +from ._nonthermal_core import workfunction_tests +from .plotnonthermal import addargs +from .plotnonthermal import main as plot diff --git a/artistools/nonthermal/__main__.py b/artistools/nonthermal/__main__.py new file mode 100644 index 000000000..3c34c1c3e --- /dev/null +++ b/artistools/nonthermal/__main__.py @@ -0,0 +1,9 @@ +from .plotnonthermal import main as plot + + +def main() -> None: + plot() + + +if __name__ == "__main__": + main() diff --git a/artistools/nonthermal/_nonthermal_core.py b/artistools/nonthermal/_nonthermal_core.py index 10e33d702..77f7bfa67 100755 --- a/artistools/nonthermal/_nonthermal_core.py +++ b/artistools/nonthermal/_nonthermal_core.py @@ -4,19 +4,14 @@ from collections import namedtuple from math import atan from pathlib import Path +from typing import Union import matplotlib.pyplot as plt import numpy as np import pandas as pd import artistools as at -import artistools.estimators -import artistools.nltepops -import artistools.nonthermal - -# import matplotlib.ticker as ticker -# import numba -# from numpy import arctan as atan +from artistools.configuration import get_config # cgs units to match artis EV = 1.6021772e-12 # in erg @@ -86,12 +81,12 @@ def read_binding_energies(modelpath: str = ".") -> np.ndarray: collionfilename = at.firstexisting( [ os.path.join(modelpath, "binding_energies.txt"), - os.path.join(at.get_config()["path_artistools_dir"], "data", "binding_energies.txt"), + os.path.join(get_config()["path_artistools_dir"], "data", "binding_energies.txt"), ] ) - with open(collionfilename, "r") as f: - nt_shells, n_z_binding = [int(x) for x in f.readline().split()] + with open(collionfilename) as f: + nt_shells, n_z_binding = (int(x) for x in f.readline().split()) electron_binding = np.zeros((n_z_binding, nt_shells)) for i in range(n_z_binding): @@ -107,7 +102,7 @@ def get_electronoccupancy(atomic_number: int, ion_stage: int, nt_shells: int) -> ioncharge = ion_stage - 1 nbound = atomic_number - ioncharge # number of bound electrons - for electron_loop in range(nbound): + for _electron_loop in range(nbound): if q[0] < 2: # K 1s q[0] += 1 elif q[1] < 2: # L1 2s @@ -253,16 +248,13 @@ def get_lotz_xs_ionisation(atomic_number, ion_stage, electron_binding, ionpot_ev assert electron_loop == 8 # print("Z = %d, ion_stage = %d\n", get_element(element), get_ionstage(element, ion)); - if use2 < use3: - p = use3 - else: - p = use2 + p = use3 if use2 < use3 else use2 if 0.5 * beta**2 * ME * CLIGHT**2 > p: part_sigma += ( electronsinshell / p - * ((math.log(beta**2 * ME * CLIGHT**2 / 2.0 / p) - math.log10(1 - beta**2) - beta**2)) + * (math.log(beta**2 * ME * CLIGHT**2 / 2.0 / p) - math.log10(1 - beta**2) - beta**2) ) Aconst = 1.33e-14 * EV * EV @@ -336,9 +328,9 @@ def get_J(Z, ionstage, ionpot_ev): if ionstage == 1: if Z == 2: # He I return 15.8 - elif Z == 10: # Ne I + if Z == 10: # Ne I return 24.2 - elif Z == 18: # Ar I + if Z == 18: # Ar I return 10.0 return 0.6 * ionpot_ev @@ -370,7 +362,7 @@ def get_arxs_array_shell(arr_enev, shell): def get_arxs_array_ion(arr_enev, dfcollion, Z, ionstage): ar_xs_array = np.zeros(len(arr_enev)) dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage") - for index, shell in dfcollion_thision.iterrows(): + for _index, shell in dfcollion_thision.iterrows(): ar_xs_array += get_arxs_array_shell(arr_enev, shell) return ar_xs_array @@ -393,7 +385,7 @@ def get_xs_excitation(en_ev, row): return constantfactor * (en_ev * EV) ** -2 - elif not row.forbidden: + if not row.forbidden: nu_trans = epsilon_trans / H g = row.upper_g / row.lower_g fij = g * ME * pow(CLIGHT, 3) / (8 * pow(QE * nu_trans * math.pi, 2)) * row.A @@ -466,17 +458,20 @@ def get_xs_excitation_vector(engrid, row): return xs_excitation_vec -def read_colliondata(collionfilename="collion.txt", modelpath=at.get_config()["path_datadir"]): +def read_colliondata(collionfilename="collion.txt", modelpath: Union[None, str, Path] = None): + if modelpath is None: + modelpath = get_config()["path_datadir"] + collionrow = namedtuple("collionrow", ["Z", "nelec", "n", "l", "ionpot_ev", "A", "B", "C", "D"]) nrows = -1 - with open(Path(modelpath, collionfilename), "r") as collionfile: + with open(Path(modelpath, collionfilename)) as collionfile: nrows = int(collionfile.readline().strip()) # print(f'Collionfile: expecting {nrows} rows') dfcollion = pd.read_csv(collionfile, delim_whitespace=True, header=None, names=collionrow._fields) - # assert len(dfcollion) == nrows # artis enforces this, but last 10 rows were not inportant anyway (high ionized Ni) - dfcollion.eval("ionstage = Z - nelec + 1", inplace=True) + # assert len(dfcollion) == nrows # artis enforces this, but last 10 rows were not inportant anyway (high ionized Ni) + dfcollion = dfcollion.eval("ionstage = Z - nelec + 1") return dfcollion @@ -506,10 +501,11 @@ def get_energyindex_lteq(en_ev, engrid): if index < 0: return 0 - elif index > len(engrid) - 1: + + if index > len(engrid) - 1: return len(engrid) - 1 - else: - return index + + return index def get_energyindex_gteq(en_ev, engrid): @@ -520,10 +516,11 @@ def get_energyindex_gteq(en_ev, engrid): if index < 0: return 0 - elif index > len(engrid) - 1: + + if index > len(engrid) - 1: return len(engrid) - 1 - else: - return index + + return index def calculate_N_e(energy_ev, engrid, ions, ionpopdict, dfcollion, yvec, dftransitions, noexcitation): @@ -541,23 +538,22 @@ def calculate_N_e(energy_ev, engrid, ions, ionpopdict, dfcollion, yvec, dftransi N_e_ion = 0.0 nnion = ionpopdict[(Z, ionstage)] - if not noexcitation: - if (Z, ionstage) in dftransitions: - for _, row in dftransitions[(Z, ionstage)].iterrows(): - nnlevel = row.lower_pop - epsilon_trans_ev = row.epsilon_trans_ev - if energy_ev + epsilon_trans_ev >= engrid[0]: - i = get_energyindex_lteq(en_ev=energy_ev + epsilon_trans_ev, engrid=engrid) - N_e_ion += (nnlevel / nnion) * yvec[i] * get_xs_excitation(engrid[i], row) - # enbelow = engrid[i] - # enabove = engrid[i + 1] - # x = (energy_ev - enbelow) / (enabove - enbelow) - # yvecinterp = (1 - x) * yvec[i] + x * yvec[i + 1] - # N_e_ion += (nnlevel / nnion) * yvecinterp * get_xs_excitation(energy_ev + epsilon_trans_ev, row) + if not noexcitation and (Z, ionstage) in dftransitions: + for _, row in dftransitions[(Z, ionstage)].iterrows(): + nnlevel = row.lower_pop + epsilon_trans_ev = row.epsilon_trans_ev + if energy_ev + epsilon_trans_ev >= engrid[0]: + i = get_energyindex_lteq(en_ev=energy_ev + epsilon_trans_ev, engrid=engrid) + N_e_ion += (nnlevel / nnion) * yvec[i] * get_xs_excitation(engrid[i], row) + # enbelow = engrid[i] + # enabove = engrid[i + 1] + # x = (energy_ev - enbelow) / (enabove - enbelow) + # yvecinterp = (1 - x) * yvec[i] + x * yvec[i + 1] + # N_e_ion += (nnlevel / nnion) * yvecinterp * get_xs_excitation(energy_ev + epsilon_trans_ev, row) dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage", inplace=False) - for index, shell in dfcollion_thision.iterrows(): + for _index, shell in dfcollion_thision.iterrows(): ionpot_ev = shell.ionpot_ev enlambda = min(engrid[-1] - energy_ev, energy_ev + ionpot_ev) @@ -713,10 +709,11 @@ def sfmatrix_add_ionization_shell(engrid, nnion, shell, sfmatrix): if epsilon_lowers1[j - i] <= epsilon_uppers[j]: sfmatrix[i, j] += prefactors[j] * (int_eps_uppers[j] - int_eps_lowers1[j - i]) - if 2 * en + ionpot_ev < engrid[-1] + (engrid[1] - engrid[0]): - secondintegralstartindex = get_energyindex_lteq(2 * en + ionpot_ev, engrid) - else: - secondintegralstartindex = npts + 1 + secondintegralstartindex = ( + get_energyindex_lteq(2 * en + ionpot_ev, engrid) + if 2 * en + ionpot_ev < engrid[-1] + (engrid[1] - engrid[0]) + else npts + 1 + ) # endash ranges from 2 * en + ionpot_ev to SF_EMAX # at each endash, the integral in epsilon ranges from @@ -739,10 +736,7 @@ def differentialsfmatrix_add_ionization_shell(engrid, nnion, shell, sfmatrix): ar_xs_array = at.nonthermal.get_arxs_array_shell(engrid, shell) - if ionpot_ev <= engrid[0]: - xsstartindex = 0 - else: - xsstartindex = get_energyindex_lteq(en_ev=ionpot_ev, engrid=engrid) + xsstartindex = 0 if ionpot_ev <= engrid[0] else get_energyindex_lteq(en_ev=ionpot_ev, engrid=engrid) oneoveratangrid = 1.0 / np.arctan((engrid - ionpot_ev) / 2.0 / J) @@ -870,7 +864,7 @@ def solve_spencerfano_differentialform( dfcollion_thision = dfcollion.query("Z == @Z and ionstage == @ionstage", inplace=False) # print(dfcollion_thision) - for index, shell in dfcollion_thision.iterrows(): + for _index, shell in dfcollion_thision.iterrows(): # assert shell.ionpot_ev >= engrid[0] if shell.ionpot_ev < engrid[0]: print(f" WARNING: first energy point at {engrid[0]} eV is above shell ionpot {shell.ionpot_ev} eV") @@ -940,7 +934,7 @@ def analyse_ntspectrum( frac_ionization_ion[(Z, ionstage)] = 0.0 # integralgamma = 0. eta_over_ionpot_sum = 0.0 - for index, shell in dfcollion_thision.iterrows(): + for _index, shell in dfcollion_thision.iterrows(): ar_xs_array = at.nonthermal.get_arxs_array_shell(engrid, shell) frac_ionization_shell = ( @@ -1110,7 +1104,7 @@ def get_fij_ln_en_ionisation(emax_ev, J, shell): e_p_lower = shell.ionpot_ev e_p_upper = emax_ev delta_e_p = (e_p_upper - e_p_lower) / npts - sum = 0.0 + sumval = 0.0 for i in range(npts): e_p = e_p_lower + i * delta_e_p print(i, e_p) @@ -1118,9 +1112,9 @@ def get_fij_ln_en_ionisation(emax_ev, J, shell): sigma = at.nonthermal.ar_xs(e_p, shell.ionpot_ev, shell.A, shell.B, shell.C, shell.D) eps_avg = get_epsilon_avg(e_p, J, shell.ionpot_ev) if eps_avg > 0: - sum += ME * CLIGHT / math.pi / (QE**2) / H * sigma * math.log(eps_avg) * delta_e_p + sumval += ME * CLIGHT / math.pi / (QE**2) / H * sigma * math.log(eps_avg) * delta_e_p - return sum + return sumval def e_s_test(ax, ionpot_ev, J, arr_en_ev, shellstr, color): @@ -1173,7 +1167,7 @@ def e_s_test(ax, ionpot_ev, J, arr_en_ev, shellstr, color): # ax.vlines(ionpot_ev, ymin=0., ymax=max(prob), color=color) -def get_epsilon_avg(e_p, J, ionpot_ev, quiet=True): +def get_epsilon_avg(e_p: float, J: float, ionpot_ev: float, quiet: bool = True) -> float: # average energy loss of the primary electron per ionisation in eV npts = 1000000 @@ -1234,15 +1228,17 @@ def calculate_Latom_excitation(ions, ionpopdict, nntot, en_ev, adata, T_exc=5000 k_b = 8.617333262145179e-05 # eV / K energy_boltzfac_sum = ion.levels.eval("energy_ev * g * exp(- energy_ev / @k_b / @T_exc)").sum() - populations = ion.levels.eval("g * exp(- energy_ev / @k_b / @T_exc)").values / energy_boltzfac_sum + populations = ion.levels.eval("g * exp(- energy_ev / @k_b / @T_exc)").to_numpy() / energy_boltzfac_sum - dftransitions_ion.eval( - "epsilon_trans_ev = @ion.levels.loc[upper].energy_ev.values - @ion.levels.loc[lower].energy_ev.values", - inplace=True, + dftransitions_ion = dftransitions_ion.eval( + ( + "epsilon_trans_ev = @ion.levels.loc[upper].energy_ev.to_numpy() -" + " @ion.levels.loc[lower].energy_ev.to_numpy()" + ), ) - dftransitions_ion.eval("upper_g = @ion.levels.loc[upper].g.values", inplace=True) - dftransitions_ion.eval("lower_g = @ion.levels.loc[lower].g.values", inplace=True) + dftransitions_ion = dftransitions_ion.eval("upper_g = @ion.levels.loc[upper].g.to_numpy()") + dftransitions_ion = dftransitions_ion.eval("lower_g = @ion.levels.loc[lower].g.to_numpy()") for _, row in dftransitions_ion.iterrows(): nnlevel = populations[row.lower] * nnion @@ -1304,7 +1300,7 @@ def workfunction_tests(modelpath, args): } ions = [] - for key in ionpopdict.keys(): + for key in ionpopdict: # keep only the ion populations, not element or total populations if isinstance(key, tuple) and len(key) == 2: ions.append(key) @@ -1339,7 +1335,6 @@ def workfunction_tests(modelpath, args): start=math.log10(en_min_ev), stop=math.log10(en_max_ev), base=10, num=args.npts, endpoint=True ) - global Psecondary_e_s_max Psecondary_e_s_max = arr_en_ev[2] print(f"Psecondary_e_s_max: {Psecondary_e_s_max}") diff --git a/artistools/nonthermal/plotnonthermal.py b/artistools/nonthermal/plotnonthermal.py index dcdf372f2..989491e34 100755 --- a/artistools/nonthermal/plotnonthermal.py +++ b/artistools/nonthermal/plotnonthermal.py @@ -13,8 +13,6 @@ import artistools as at -# import matplotlib.ticker as ticker - DEFAULTSPECPATH = "../example_run/spec.out" defaultoutputfile = "plotnonthermal_cell{0:03d}_timestep{1:03d}.pdf" @@ -27,7 +25,7 @@ def read_files(modelpath, timestep=-1, modelgridindex=-1): mpiranklist = at.get_mpiranklist(modelpath, modelgridindex=modelgridindex) for folderpath in at.get_runfolders(modelpath, timestep=timestep): for mpirank in mpiranklist: - filepath = at.firstexisting(f"nonthermalspec_{mpirank:04d}.out", path=folderpath, tryzipped=True) + filepath = at.firstexisting(f"nonthermalspec_{mpirank:04d}.out", folder=folderpath, tryzipped=True) if modelgridindex > -1: filesize = Path(filepath).stat().st_size / 1024 / 1024 @@ -37,16 +35,16 @@ def read_files(modelpath, timestep=-1, modelgridindex=-1): # radfielddata_thisfile[['modelgridindex', 'timestep']].apply(pd.to_numeric) if timestep >= 0: - nonthermaldata_thisfile.query("timestep==@timestep", inplace=True) + nonthermaldata_thisfile = nonthermaldata_thisfile.query("timestep==@timestep") if modelgridindex >= 0: - nonthermaldata_thisfile.query("modelgridindex==@modelgridindex", inplace=True) + nonthermaldata_thisfile = nonthermaldata_thisfile.query("modelgridindex==@modelgridindex") if not nonthermaldata_thisfile.empty: if timestep >= 0 and modelgridindex >= 0: return nonthermaldata_thisfile - else: - nonthermaldata = nonthermaldata.append(nonthermaldata_thisfile.copy(), ignore_index=True) + + nonthermaldata = nonthermaldata.append(nonthermaldata_thisfile.copy(), ignore_index=True) return nonthermaldata @@ -83,8 +81,8 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata print(f"Deposition: {total_depev:.1f} [eV/cm3/s]") - arr_enev = nonthermaldata["energy_ev"].values - arr_y = nonthermaldata["y"].values + arr_enev = nonthermaldata["energy_ev"].to_numpy() + arr_y = nonthermaldata["y"].to_numpy() frac_ionisation = 0.0 @@ -115,7 +113,7 @@ def plot_contributions(axis, modelpath, timestep, modelgridindex, nonthermaldata arr_ionisation_ion = np.zeros(len(arr_enev), dtype=float) frac_ionisation_ion = 0.0 - for index, row in dfcollion_thision.iterrows(): + for _index, row in dfcollion_thision.iterrows(): arr_xs = at.nonthermal.get_arxs_array_shell(arr_enev, row) arr_ionisation_shell = ionpop * arr_y * arr_xs * row.ionpot_ev / total_depev arr_ionisation_ion += arr_ionisation_shell @@ -167,27 +165,25 @@ def make_plot(modelpaths, args): if args.kf1992spec: kf92spec = pd.read_csv(Path(modelpaths[0], "KF1992spec-fig1.txt"), header=None, names=["e_kev", "log10_y"]) kf92spec["energy_ev"] = kf92spec["e_kev"] * 1000.0 - kf92spec.eval("y = 10 ** log10_y", inplace=True) + kf92spec = kf92spec.eval("y = 10 ** log10_y") axes[0].plot( kf92spec["energy_ev"], kf92spec["log10_y"], linewidth=2.0, color="red", label="Kozma & Fransson (1992)" ) for index, modelpath in enumerate(modelpaths): modelname = at.get_model_name(modelpath) - if args.velocity >= 0.0: - modelgridindex = at.inputmodel.get_mgi_of_velocity_kms(modelpath, args.velocity) - else: - modelgridindex = args.modelgridindex + modelgridindex = ( + at.inputmodel.get_mgi_of_velocity_kms(modelpath, args.velocity) + if args.velocity >= 0.0 + else args.modelgridindex + ) - if args.timedays: - timestep = at.get_timestep_of_timedays(modelpath, args.timedays) - else: - timestep = args.timestep + timestep = at.get_timestep_of_timedays(modelpath, args.timedays) if args.timedays else args.timestep nonthermaldata = read_files(modelpath=Path(modelpath), modelgridindex=modelgridindex, timestep=timestep) if args.xmin: - nonthermaldata.query("energy_ev >= @args.xmin", inplace=True) + nonthermaldata = nonthermaldata.query("energy_ev >= @args.xmin") if nonthermaldata.empty: print(f"No data for timestep {timestep:d}") diff --git a/artistools/nonthermal/solvespencerfanocmd.py b/artistools/nonthermal/solvespencerfanocmd.py old mode 100644 new mode 100755 index 3368cc636..d21f4ed26 --- a/artistools/nonthermal/solvespencerfanocmd.py +++ b/artistools/nonthermal/solvespencerfanocmd.py @@ -3,6 +3,7 @@ import multiprocessing import sys from pathlib import Path +from typing import Union import matplotlib.pyplot as plt import numpy as np @@ -10,31 +11,24 @@ import pynonthermal as pynt import artistools as at -import artistools.estimators -import artistools.nltepops -import artistools.nonthermal - -# import numba -# from numpy import arctan as atan - minionfraction = 0.0 # minimum number fraction of the total population to include in SF solution defaultoutputfile = "spencerfano_cell{cell:03d}_ts{timestep:02d}_{timedays:.0f}d.pdf" -def make_ntstats_plot(ntstatfile): +def make_ntstats_plot(ntstatfile: Union[str, Path]) -> None: fig, ax = plt.subplots( nrows=1, ncols=1, sharex=True, figsize=(4, 3), tight_layout={"pad": 0.5, "w_pad": 0.3, "h_pad": 0.3} ) dfstats = pd.read_csv(ntstatfile, delim_whitespace=True, escapechar="#") - dfstats.fillna(0, inplace=True) + dfstats = dfstats.fillna(0) norm_frac_sum = False if norm_frac_sum: # scale up (or down) ionisation, excitation, and heating to force frac_sum = 1.0 - dfstats.eval("frac_sum = frac_ionization + frac_excitation + frac_heating", inplace=True) + dfstats = dfstats.eval("frac_sum = frac_ionization + frac_excitation + frac_heating") norm_factors = 1.0 / dfstats["frac_sum"] else: norm_factors = 1.0 @@ -45,10 +39,10 @@ def make_ntstats_plot(ntstatfile): xarr = np.log10(dfstats.x_e) ax.plot(xarr, dfstats.frac_ionization * norm_factors, label="Ionisation") - if not max(dfstats.frac_excitation) == 0.0: + if max(dfstats.frac_excitation) > 0.0: ax.plot(xarr, dfstats.frac_excitation * norm_factors, label="Excitation") ax.plot(xarr, dfstats.frac_heating * norm_factors, label="Heating") - ioncols = [col for col in dfstats.columns.values if col.startswith("frac_ionization_")] + ioncols = [col for col in dfstats.columns.to_numpy() if col.startswith("frac_ionization_")] for ioncol in ioncols: ion = ioncol.replace("frac_ionization_", "") ax.plot(xarr, dfstats[ioncol] * norm_factors, label=f"{ion} ionisation") @@ -160,7 +154,7 @@ def main(args=None, argsraw=None, **kwargs): if args.plotstats: make_ntstats_plot(args.plotstats) - return + return None # global at.nonthermal.experiment_use_Latom_in_spencerfano at.nonthermal.experiment_use_Latom_in_spencerfano = args.atomlossrate @@ -314,7 +308,7 @@ def main(args=None, argsraw=None, **kwargs): ionpopdict[(compelement_atomicnumber, 2)] = nntot * x_e ions = [] - for key in ionpopdict.keys(): + for key in ionpopdict: # keep only the ion populations, not element or total populations if isinstance(key, tuple) and len(key) == 2 and ionpopdict[key] / nntot >= minionfraction: ions.append(key) @@ -365,15 +359,14 @@ def main(args=None, argsraw=None, **kwargs): ) for atomic_number, ionstage in ions: nnion = ionpopdict[(atomic_number, ionstage)] - if nnion > 0.0: - frac_ionis_ion = sf.get_frac_ionisation_ion(atomic_number, ionstage) - else: - frac_ionis_ion = 0.0 + frac_ionis_ion = sf.get_frac_ionisation_ion(atomic_number, ionstage) if nnion > 0.0 else 0.0 strlineout += f" {frac_ionis_ion:.4f}" fstat.write(strlineout + "\n") if args.ostat: make_ntstats_plot(args.ostat) + return None + return None if __name__ == "__main__": diff --git a/artistools/packets/__init__.py b/artistools/packets/__init__.py index 8d1a56cf2..900f30f89 100644 --- a/artistools/packets/__init__.py +++ b/artistools/packets/__init__.py @@ -1,539 +1,14 @@ -#!/usr/bin/env python3 -import gzip -import math -import sys -from collections.abc import Sequence -from functools import lru_cache -from pathlib import Path -from typing import Optional -from typing import Union - -import numpy as np -import pandas as pd - import artistools as at - -# import multiprocessing -# import matplotlib.patches as mpatches -# from collections import namedtuple - -CLIGHT = 2.99792458e10 -DAY = 86400 - -types = { - 10: "TYPE_GAMMA", - 11: "TYPE_RPKT", - 20: "TYPE_NTLEPTON", - 32: "TYPE_ESCAPE", -} - -type_ids = dict((v, k) for k, v in types.items()) - - -@lru_cache(maxsize=16) -def get_column_names_artiscode(modelpath: Union[str, Path]) -> Optional[list[str]]: - modelpath = Path(modelpath) - if Path(modelpath, "artis").is_dir(): - print("detected artis code directory") - packet_properties = [] - inputfilename = at.firstexisting(["packet_init.cc", "packet_init.c"], path=(modelpath / "artis")) - print(f"found {inputfilename}: getting packet column names from artis code") - with open(inputfilename) as inputfile: - packet_print_lines = [line.split(",") for line in inputfile if "fprintf(packets_file," in line] - for line in packet_print_lines: - for element in line: - if "pkt[i]." in element: - packet_properties.append(element) - - for i, element in enumerate(packet_properties): - packet_properties[i] = element.split(".")[1].split(")")[0] - - columns = packet_properties - replacements_dict = { - "type": "type_id", - "pos[0]": "posx", - "pos[1]": "posy", - "pos[2]": "posz", - "dir[0]": "dirx", - "dir[1]": "diry", - "dir[2]": "dirz", - "escape_type": "escape_type_id", - "em_pos[0]": "em_posx", - "em_pos[1]": "em_posy", - "em_pos[2]": "em_posz", - "absorptiontype": "absorption_type", - "absorptionfreq": "absorption_freq", - "absorptiondir[0]": "absorptiondirx", - "absorptiondir[1]": "absorptiondiry", - "absorptiondir[2]": "absorptiondirz", - "stokes[0]": "stokes1", - "stokes[1]": "stokes2", - "stokes[2]": "stokes3", - "pol_dir[0]": "pol_dirx", - "pol_dir[1]": "pol_diry", - "pol_dir[2]": "pol_dirz", - "trueemissionvelocity": "true_emission_velocity", - } - - for i, column_name in enumerate(columns): - if column_name in replacements_dict: - columns[i] = replacements_dict[column_name] - - return columns - - return None - - -def add_derived_columns( - dfpackets: pd.DataFrame, - modelpath: Path, - colnames: Sequence[str], - allnonemptymgilist: Optional[Sequence[int]] = None, -) -> pd.DataFrame: - cm_to_km = 1e-5 - day_in_s = 86400 - if dfpackets.empty: - return dfpackets - - colnames = at.makelist(colnames) - - def em_modelgridindex(packet) -> Union[int, float]: - return at.inputmodel.get_mgi_of_velocity_kms( - modelpath, packet.emission_velocity * cm_to_km, mgilist=allnonemptymgilist - ) - - def emtrue_modelgridindex(packet) -> Union[int, float]: - return at.inputmodel.get_mgi_of_velocity_kms( - modelpath, packet.true_emission_velocity * cm_to_km, mgilist=allnonemptymgilist - ) - - def em_timestep(packet) -> int: - return at.get_timestep_of_timedays(modelpath, packet.em_time / day_in_s) - - def emtrue_timestep(packet) -> int: - return at.get_timestep_of_timedays(modelpath, packet.trueem_time / day_in_s) - - if "emission_velocity" in colnames: - dfpackets.eval("emission_velocity = sqrt(em_posx ** 2 + em_posy ** 2 + em_posz ** 2) / em_time", inplace=True) - - dfpackets.eval("em_velx = em_posx / em_time", inplace=True) - dfpackets.eval("em_vely = em_posy / em_time", inplace=True) - dfpackets.eval("em_velz = em_posz / em_time", inplace=True) - - if "em_modelgridindex" in colnames: - if "emission_velocity" not in dfpackets.columns: - dfpackets = add_derived_columns( - dfpackets, modelpath, ["emission_velocity"], allnonemptymgilist=allnonemptymgilist - ) - dfpackets["em_modelgridindex"] = dfpackets.apply(em_modelgridindex, axis=1) - - if "emtrue_modelgridindex" in colnames: - dfpackets["emtrue_modelgridindex"] = dfpackets.apply(emtrue_modelgridindex, axis=1) - - if "em_timestep" in colnames: - dfpackets["em_timestep"] = dfpackets.apply(em_timestep, axis=1) - - if "angle_bin" in colnames: - dfpackets = get_escaping_packet_angle_bin(modelpath, dfpackets) - - return dfpackets - - -def readfile_text(packetsfile: Union[Path, str], modelpath: Path = Path(".")) -> pd.DataFrame: - usecols_nodata = None # print a warning for missing columns if the source code columns can't be read - - skiprows: int = 0 - column_names: Optional[list[str]] = None - try: - fpackets = at.zopen(packetsfile, "rt") - - datastartpos = fpackets.tell() # will be updated if this was actually the start of a header - firstline = fpackets.readline() - - if firstline.lstrip().startswith("#"): - column_names = firstline.lstrip("#").split() - assert column_names is not None - column_names.append("ignore") - # get the column count from the first data line to check header matched - datastartpos = fpackets.tell() - dataline = fpackets.readline() - inputcolumncount = len(dataline.split()) - skiprows = 1 - else: - inputcolumncount = len(firstline.split()) - - fpackets.seek(datastartpos) # go to first data line - - except gzip.BadGzipFile: - print(f"\nBad Gzip File: {packetsfile}") - raise gzip.BadGzipFile - - try: - dfpackets = pd.read_csv( - fpackets, - sep=" ", - header=None, - skiprows=skiprows, - names=column_names, - skip_blank_lines=True, - engine="pyarrow", - ) - except ImportError: - dfpackets = pd.read_csv( - fpackets, - sep=" ", - header=None, - skiprows=skiprows, - names=column_names, - skip_blank_lines=True, - ) - - # import datatable as dt - # dsk_dfpackets = dt.fread(packetsfile) - # dfpackets = dsk_dfpackets.to_pandas() - - # space at the end of line made an extra column of Nones - if dfpackets[dfpackets.columns[-1]].isnull().all(): - dfpackets.drop(labels=dfpackets.columns[-1], axis=1, inplace=True) - - if hasattr(dfpackets.columns[0], "startswith") and dfpackets.columns[0].startswith("#"): - dfpackets.rename(columns={dfpackets.columns[0]: dfpackets.columns[0].lstrip("#")}, inplace=True) - - elif dfpackets.columns[0] in ["C0", 0]: - inputcolumncount = len(dfpackets.columns) - column_names = get_column_names_artiscode(modelpath) - if column_names: # found them in the artis code files - assert len(column_names) == inputcolumncount - - else: # infer from column positions - # new artis added extra columns to the end of this list, but they may be absent in older versions - # the packets file may have a truncated set of columns, but we assume that they - # are only truncated, i.e. the columns with the same index have the same meaning - columns_full = [ - "number", - "where", - "type_id", - "posx", - "posy", - "posz", - "dirx", - "diry", - "dirz", - "last_cross", - "tdecay", - "e_cmf", - "e_rf", - "nu_cmf", - "nu_rf", - "escape_type_id", - "escape_time", - "scat_count", - "next_trans", - "interactions", - "last_event", - "emissiontype", - "trueemissiontype", - "em_posx", - "em_posy", - "em_posz", - "absorption_type", - "absorption_freq", - "nscatterings", - "em_time", - "absorptiondirx", - "absorptiondiry", - "absorptiondirz", - "stokes1", - "stokes2", - "stokes3", - "pol_dirx", - "pol_diry", - "pol_dirz", - "originated_from_positron", - "true_emission_velocity", - "trueem_time", - "pellet_nucindex", - ] - - assert len(columns_full) >= inputcolumncount - usecols_nodata = [n for n in columns_full if columns_full.index(n) >= inputcolumncount] - column_names = columns_full[:inputcolumncount] - - dfpackets.columns = column_names - - # except Exception as ex: - # print(f'Problem with file {packetsfile}') - # print(f'ERROR: {ex}') - # sys.exit(1) - - if usecols_nodata: - print(f"WARNING: no data in packets file for columns: {usecols_nodata}") - for col in usecols_nodata: - dfpackets[col] = float("NaN") - - return dfpackets - - -@at.diskcache(savezipped=True) -def readfile( - packetsfile: Union[Path, str], type: Optional[str] = None, escape_type: Optional[str] = None -) -> pd.DataFrame: - """Read a packet file into a pandas DataFrame.""" - packetsfile = Path(packetsfile) - - if packetsfile.suffixes == [".out", ".parquet"]: - dfpackets = pd.read_parquet(packetsfile) - elif packetsfile.suffixes == [".out", ".feather"]: - dfpackets = pd.read_feather(packetsfile) - elif packetsfile.suffixes in [[".out"], [".out", ".gz"], [".out", ".xz"]]: - dfpackets = readfile_text(packetsfile) - # dfpackets.to_parquet(at.stripallsuffixes(packetsfile).with_suffix('.out.parquet'), - # compression='brotli', compression_level=99) - else: - print("ERROR") - sys.exit(1) - filesize = Path(packetsfile).stat().st_size / 1024 / 1024 - print(f"Reading {packetsfile} ({filesize:.1f} MiB)", end="") - - print(f" ({len(dfpackets):.1e} packets", end="") - - if escape_type is not None and escape_type != "" and escape_type != "ALL": - assert type is None or type == "TYPE_ESCAPE" - dfpackets.query( - f'type_id == {type_ids["TYPE_ESCAPE"]} and escape_type_id == {type_ids[escape_type]}', inplace=True - ) - print(f", {len(dfpackets)} escaped as {escape_type})") - elif type is not None and type != "ALL" and type != "": - dfpackets.query(f"type_id == {type_ids[type]}", inplace=True) - print(f", {len(dfpackets)} with type {type})") - else: - print(")") - - # dfpackets['type'] = dfpackets['type_id'].map(lambda x: types.get(x, x)) - # dfpackets['escape_type'] = dfpackets['escape_type_id'].map(lambda x: types.get(x, x)) - - # # neglect light travel time correction - # dfpackets.eval("t_arrive_d = escape_time / 86400", inplace=True) - - dfpackets.eval( - "t_arrive_d = (escape_time - (posx * dirx + posy * diry + posz * dirz) / 29979245800) / 86400", inplace=True - ) - - return dfpackets - - -@lru_cache(maxsize=16) -def get_packetsfilepaths(modelpath: Path, maxpacketfiles: Optional[int] = None) -> list[Path]: - def preferred_alternative(f: Path) -> bool: - f_nosuffixes = at.stripallsuffixes(f) - - suffix_priority = [[".out", ".gz"], [".out", ".xz"], [".out", ".feather"], [".out", ".parquet"]] - if f.suffixes in suffix_priority: - startindex = suffix_priority.index(f.suffixes) + 1 - else: - startindex = 0 - - if any(f_nosuffixes.with_suffix("".join(s)).is_file() for s in suffix_priority[startindex:]): - return True - return False - - packetsfiles = sorted( - list(Path(modelpath).glob("packets00_*.out*")) + list(Path(modelpath, "packets").glob("packets00_*.out*")) - ) - - # strip out duplicates in the case that some are stored as binary and some are text files - packetsfiles = [f for f in packetsfiles if not preferred_alternative(f)] - - if maxpacketfiles is not None and maxpacketfiles > 0 and len(packetsfiles) > maxpacketfiles: - print(f"Using only the first {maxpacketfiles} of {len(packetsfiles)} packets files") - packetsfiles = packetsfiles[:maxpacketfiles] - - return packetsfiles - - -def get_directionbin(dirx: float, diry: float, dirz: float, nphibins: int, syn_dir: Sequence[float]) -> int: - dirmag = np.sqrt(dirx**2 + diry**2 + dirz**2) - pkt_dir = [dirx / dirmag, diry / dirmag, dirz / dirmag] - costheta = np.dot(pkt_dir, syn_dir) - thetabin = int((costheta + 1.0) * nphibins / 2.0) - vec1 = vec2 = vec3 = np.array([0.0, 0.0, 0.0]) - xhat = np.array([1.0, 0.0, 0.0]) - vec1 = np.cross(pkt_dir, syn_dir) - vec2 = np.cross(xhat, syn_dir) - cosphi = np.dot(vec1, vec2) / at.vec_len(vec1) / at.vec_len(vec2) - - vec3 = np.cross(vec2, syn_dir) - testphi = np.dot(vec1, vec3) - - if testphi > 0: - phibin = int(math.acos(cosphi) / 2.0 / np.pi * nphibins) - else: - phibin = int((math.acos(cosphi) + np.pi) / 2.0 / np.pi * nphibins) - - na = (thetabin * nphibins) + phibin - return na - - -def get_escaping_packet_angle_bin(modelpath: Union[Path, str], dfpackets: pd.DataFrame) -> pd.DataFrame: - nphibins = at.get_viewingdirection_phibincount() - - syn_dir = at.get_syn_dir(Path(modelpath)) - - get_all_directionbins = np.vectorize(get_directionbin, excluded=["nphibins", "syn_dir"]) - dfpackets["angle_bin"] = get_all_directionbins( - dfpackets["dirx"], dfpackets["diry"], dfpackets["dirz"], nphibins=nphibins, syn_dir=syn_dir - ) - assert np.all(dfpackets["angle_bin"] < at.get_viewingdirectionbincount()) - - return dfpackets - - -def make_3d_histogram_from_packets(modelpath, timestep_min, timestep_max=None, em_time=True): - if timestep_max is None: - timestep_max = timestep_min - modeldata, _, vmax_cms = at.inputmodel.get_modeldata_tuple(modelpath) - - timeminarray = at.get_timestep_times_float(modelpath=modelpath, loc="start") - timedeltaarray = at.get_timestep_times_float(modelpath=modelpath, loc="delta") - timemaxarray = at.get_timestep_times_float(modelpath=modelpath, loc="end") - - # timestep = 63 # 82 73 #63 #54 46 #27 - # print([(ts, time) for ts, time in enumerate(timeminarray)]) - if em_time: - print("Binning by packet emission time") - else: - print("Binning by packet arrival time") - - packetsfiles = at.packets.get_packetsfilepaths(modelpath) - - emission_position3d = [[], [], []] - e_rf = [] - e_cmf = [] - - for npacketfile in range(0, len(packetsfiles)): - # for npacketfile in range(0, 1): - dfpackets = at.packets.readfile(packetsfiles[npacketfile]) - at.packets.add_derived_columns(dfpackets, modelpath, ["emission_velocity"]) - dfpackets = dfpackets.dropna(subset=["emission_velocity"]) # drop rows where emission_vel is NaN - - only_packets_0_scatters = False - if only_packets_0_scatters: - print("Only using packets with 0 scatters") - # print(dfpackets[['scat_count', 'interactions', 'nscatterings']]) - dfpackets.query("nscatterings == 0", inplace=True) - - # print(dfpackets[['emission_velocity', 'em_velx', 'em_vely', 'em_velz']]) - # select only type escape and type r-pkt (don't include gamma-rays) - dfpackets.query( - f'type_id == {type_ids["TYPE_ESCAPE"]} and escape_type_id == {type_ids["TYPE_RPKT"]}', inplace=True - ) - if em_time: - dfpackets.query("@timeminarray[@timestep_min] < em_time/@DAY < @timemaxarray[@timestep_max]", inplace=True) - else: # packet arrival time - dfpackets.query("@timeminarray[@timestep_min] < t_arrive_d < @timemaxarray[@timestep_max]", inplace=True) - - emission_position3d[0].extend([em_velx / CLIGHT for em_velx in dfpackets["em_velx"]]) - emission_position3d[1].extend([em_vely / CLIGHT for em_vely in dfpackets["em_vely"]]) - emission_position3d[2].extend([em_velz / CLIGHT for em_velz in dfpackets["em_velz"]]) - - e_rf.extend([e_rf for e_rf in dfpackets["e_rf"]]) - e_cmf.extend([e_cmf for e_cmf in dfpackets["e_cmf"]]) - - emission_position3d = np.array(emission_position3d) - weight_by_energy = True - if weight_by_energy: - e_rf = np.array(e_rf) - e_cmf = np.array(e_cmf) - # weights = e_rf - weights = e_cmf - else: - weights = None - - print(emission_position3d.shape) - print(emission_position3d[0].shape) - - # print(emission_position3d) - grid_3d, _, _, _ = make_3d_grid(modeldata, vmax_cms) - print(grid_3d.shape) - # https://stackoverflow.com/questions/49861468/binning-random-data-to-regular-3d-grid-with-unequal-axis-lengths - hist, _ = np.histogramdd(emission_position3d.T, [np.append(ax, np.inf) for ax in grid_3d], weights=weights) - # print(hist.shape) - if weight_by_energy: - # Divide binned energies by number of processes and by length of timestep - hist = ( - hist / len(packetsfiles) / (timemaxarray[timestep_max] - timeminarray[timestep_min]) - ) # timedeltaarray[timestep] # histogram weighted by energy - # - need to divide by number of processes - # and length of timestep(s) - - # # print histogram coordinates - # coords = np.nonzero(hist) - # for i, j, k in zip(*coords): - # print(f'({grid_3d[0][i]}, {grid_3d[1][j]}, {grid_3d[2][k]}): {hist[i][j][k]}') - - return hist - - -def make_3d_grid(modeldata, vmax_cms): - # modeldata, _, vmax_cms = at.inputmodel.get_modeldata_tuple(modelpath) - grid = round(len(modeldata["inputcellid"]) ** (1.0 / 3.0)) - xgrid = np.zeros(grid) - vmax = vmax_cms / CLIGHT - i = 0 - for z in range(0, grid): - for y in range(0, grid): - for x in range(0, grid): - xgrid[x] = -vmax + 2 * x * vmax / grid - i += 1 - - x, y, z = np.meshgrid(xgrid, xgrid, xgrid) - grid_3d = np.array([xgrid, xgrid, xgrid]) - # grid_Te = np.zeros((grid, grid, grid)) - # print(grid_Te.shape) - return grid_3d, x, y, z - - -def get_mean_packet_emission_velocity_per_ts( - modelpath, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT", maxpacketfiles=None, escape_angles=None -) -> pd.DataFrame: - packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles=maxpacketfiles) - nprocs_read = len(packetsfiles) - assert nprocs_read > 0 - - timearray = at.get_timestep_times_float(modelpath=modelpath, loc="mid") - arr_timedelta = at.get_timestep_times_float(modelpath=modelpath, loc="delta") - timearrayplusend = np.concatenate([timearray, [timearray[-1] + arr_timedelta[-1]]]) - - dfpackets_escape_velocity_and_arrive_time = pd.DataFrame - emission_data = pd.DataFrame( - {"t_arrive_d": timearray, "mean_emission_velocity": np.zeros_like(timearray, dtype=float)} - ) - - for i, packetsfile in enumerate(packetsfiles): - dfpackets = at.packets.readfile(packetsfile, type=packet_type, escape_type=escape_type) - at.packets.add_derived_columns(dfpackets, modelpath, ["emission_velocity"]) - if escape_angles is not None: - dfpackets = at.packets.get_escaping_packet_angle_bin(modelpath, dfpackets) - dfpackets.query("angle_bin == @escape_angles", inplace=True) - - if i == 0: # make new df - dfpackets_escape_velocity_and_arrive_time = dfpackets[["t_arrive_d", "emission_velocity"]] - else: # append to df - # dfpackets_escape_velocity_and_arrive_time = dfpackets_escape_velocity_and_arrive_time.append( - # other=dfpackets[["t_arrive_d", "emission_velocity"]], ignore_index=True - # ) - dfpackets_escape_velocity_and_arrive_time = pd.concat( - [dfpackets_escape_velocity_and_arrive_time, dfpackets[["t_arrive_d", "emission_velocity"]]], - ignore_index=True, - ) - - print(dfpackets_escape_velocity_and_arrive_time) - binned = pd.cut( - dfpackets_escape_velocity_and_arrive_time["t_arrive_d"], timearrayplusend, labels=False, include_lowest=True - ) - for binindex, emission_velocity in ( - dfpackets_escape_velocity_and_arrive_time.groupby(binned)["emission_velocity"].mean().iteritems() - ): - emission_data["mean_emission_velocity"][binindex] += emission_velocity # / 2.99792458e10 - - return emission_data +from .packets import add_derived_columns +from .packets import bin_and_sum +from .packets import bin_packet_directions +from .packets import get_directionbin +from .packets import get_mean_packet_emission_velocity_per_ts +from .packets import get_packets_pl +from .packets import get_packetsfilepaths +from .packets import make_3d_grid +from .packets import make_3d_histogram_from_packets +from .packets import readfile +from .packets import readfile_pl +from .packets import readfile_text +from .packets import type_ids diff --git a/artistools/packets/packets.py b/artistools/packets/packets.py new file mode 100644 index 000000000..a3e3b816b --- /dev/null +++ b/artistools/packets/packets.py @@ -0,0 +1,775 @@ +import calendar +import gzip +import math +import os +from collections.abc import Sequence +from functools import lru_cache +from pathlib import Path +from typing import Literal +from typing import Optional +from typing import Union + +import numpy as np +import pandas as pd +import polars as pl + +import artistools as at + +# for the parquet files +time_lastschemachange = (2023, 4, 2, 22, 13, 0) + +CLIGHT = 2.99792458e10 +DAY = 86400 + +types = { + 10: "TYPE_GAMMA", + 11: "TYPE_RPKT", + 20: "TYPE_NTLEPTON", + 32: "TYPE_ESCAPE", +} + +type_ids = {v: k for k, v in types.items()} + +EMTYPE_NOTSET = -9999000 +EMTYPE_FREEFREE = -9999999 + +# new artis added extra columns to the end of this list, but they may be absent in older versions +# the packets file may have a truncated set of columns, but we assume that they +# are only truncated, i.e. the columns with the same index have the same meaning +columns_full = [ + "number", + "where", + "type_id", + "posx", + "posy", + "posz", + "dirx", + "diry", + "dirz", + "last_cross", + "tdecay", + "e_cmf", + "e_rf", + "nu_cmf", + "nu_rf", + "escape_type_id", + "escape_time", + "scat_count", + "next_trans", + "interactions", + "last_event", + "emissiontype", + "trueemissiontype", + "em_posx", + "em_posy", + "em_posz", + "absorption_type", + "absorption_freq", + "nscatterings", + "em_time", + "absorptiondirx", + "absorptiondiry", + "absorptiondirz", + "stokes1", + "stokes2", + "stokes3", + "pol_dirx", + "pol_diry", + "pol_dirz", + "originated_from_positron", + "true_emission_velocity", + "trueem_time", + "pellet_nucindex", +] + + +@lru_cache(maxsize=16) +def get_column_names_artiscode(modelpath: Union[str, Path]) -> Optional[list[str]]: + modelpath = Path(modelpath) + if Path(modelpath, "artis").is_dir(): + print("detected artis code directory") + packet_properties = [] + inputfilename = at.firstexisting(["packet_init.cc", "packet_init.c"], folder=modelpath / "artis") + print(f"found {inputfilename}: getting packet column names from artis code") + with inputfilename.open() as inputfile: + packet_print_lines = [line.split(",") for line in inputfile if "fprintf(packets_file," in line] + for line in packet_print_lines: + for element in line: + if "pkt[i]." in element: + packet_properties.append(element) + + for i, element in enumerate(packet_properties): + packet_properties[i] = element.split(".")[1].split(")")[0] + + columns = packet_properties + replacements_dict = { + "type": "type_id", + "pos[0]": "posx", + "pos[1]": "posy", + "pos[2]": "posz", + "dir[0]": "dirx", + "dir[1]": "diry", + "dir[2]": "dirz", + "escape_type": "escape_type_id", + "em_pos[0]": "em_posx", + "em_pos[1]": "em_posy", + "em_pos[2]": "em_posz", + "absorptiontype": "absorption_type", + "absorptionfreq": "absorption_freq", + "absorptiondir[0]": "absorptiondirx", + "absorptiondir[1]": "absorptiondiry", + "absorptiondir[2]": "absorptiondirz", + "stokes[0]": "stokes1", + "stokes[1]": "stokes2", + "stokes[2]": "stokes3", + "pol_dir[0]": "pol_dirx", + "pol_dir[1]": "pol_diry", + "pol_dir[2]": "pol_dirz", + "trueemissionvelocity": "true_emission_velocity", + } + + for i, column_name in enumerate(columns): + if column_name in replacements_dict: + columns[i] = replacements_dict[column_name] + + return columns + + return None + + +def add_derived_columns( + dfpackets: pd.DataFrame, + modelpath: Path, + colnames: Sequence[str], + allnonemptymgilist: Optional[Sequence[int]] = None, +) -> pd.DataFrame: + cm_to_km = 1e-5 + day_in_s = 86400 + if dfpackets.empty: + return dfpackets + + colnames = at.makelist(colnames) + + def em_modelgridindex(packet) -> Union[int, float]: + return at.inputmodel.get_mgi_of_velocity_kms( + modelpath, packet.emission_velocity * cm_to_km, mgilist=allnonemptymgilist + ) + + def emtrue_modelgridindex(packet) -> Union[int, float]: + return at.inputmodel.get_mgi_of_velocity_kms( + modelpath, packet.true_emission_velocity * cm_to_km, mgilist=allnonemptymgilist + ) + + def em_timestep(packet) -> int: + return at.get_timestep_of_timedays(modelpath, packet.em_time / day_in_s) + + def emtrue_timestep(packet) -> int: + return at.get_timestep_of_timedays(modelpath, packet.trueem_time / day_in_s) + + if "emission_velocity" in colnames: + dfpackets = dfpackets.eval("emission_velocity = sqrt(em_posx ** 2 + em_posy ** 2 + em_posz ** 2) / em_time") + dfpackets["emission_velocity"] = ( + np.sqrt(dfpackets["em_posx"] ** 2 + dfpackets["em_posy"] ** 2 + dfpackets["em_posz"] ** 2) + / dfpackets["em_time"] + ) + + dfpackets["em_velx"] = dfpackets["em_posx"] / dfpackets["em_time"] + dfpackets["em_vely"] = dfpackets["em_posy"] / dfpackets["em_time"] + dfpackets["em_velz"] = dfpackets["em_posz"] / dfpackets["em_time"] + + if "em_modelgridindex" in colnames: + if "emission_velocity" not in dfpackets.columns: + dfpackets = add_derived_columns( + dfpackets, modelpath, ["emission_velocity"], allnonemptymgilist=allnonemptymgilist + ) + dfpackets["em_modelgridindex"] = dfpackets.apply(em_modelgridindex, axis=1) + + if "emtrue_modelgridindex" in colnames: + dfpackets["emtrue_modelgridindex"] = dfpackets.apply(emtrue_modelgridindex, axis=1) + + if "emtrue_timestep" in colnames: + dfpackets["emtrue_timestep"] = dfpackets.apply(emtrue_timestep, axis=1) + + if "em_timestep" in colnames: + dfpackets["em_timestep"] = dfpackets.apply(em_timestep, axis=1) + + if any(x in colnames for x in ["angle_bin", "dirbin", "costhetabin", "phibin"]): + dfpackets = bin_packet_directions(modelpath, dfpackets) + + return dfpackets + + +def readfile_text(packetsfile: Union[Path, str], modelpath: Path = Path(".")) -> pl.DataFrame: + """Read a packets*.out(.xz) space-separated text file into a polars DataFrame.""" + print(f"Reading {packetsfile}") + skiprows: int = 0 + column_names: Optional[list[str]] = None + try: + fpackets = at.zopen(packetsfile, mode="rb") + + datastartpos = fpackets.tell() # will be updated if this was actually the start of a header + firstline = fpackets.readline().decode() + + if firstline.lstrip().startswith("#"): + column_names = firstline.lstrip("#").split() + assert column_names is not None + + # get the column count from the first data line to check header matched + datastartpos = fpackets.tell() + dataline = fpackets.readline().decode() + inputcolumncount = len(dataline.split()) + assert inputcolumncount == len(column_names) + skiprows = 1 + else: + inputcolumncount = len(firstline.split()) + column_names = get_column_names_artiscode(modelpath) + if column_names: # found them in the artis code files + assert len(column_names) == inputcolumncount + + else: # infer from column positions + assert len(columns_full) >= inputcolumncount + column_names = columns_full[:inputcolumncount] + + fpackets.seek(datastartpos) # go to first data line + + except gzip.BadGzipFile: + print(f"\nBad Gzip File: {packetsfile}") + raise + + try: + dfpackets = pl.read_csv( + fpackets, + separator=" ", + has_header=False, + new_columns=column_names, + infer_schema_length=10000, + ) + + except Exception: + print(f"Error occured in file {packetsfile}") + raise + + dfpackets = dfpackets.drop(["next_trans", "last_cross"]) + + # drop last column of nulls (caused by trailing space on each line) + if dfpackets[dfpackets.columns[-1]].is_null().all(): + dfpackets = dfpackets.drop(dfpackets.columns[-1]) + + if "true_emission_velocity" in dfpackets.columns: + dfpackets = dfpackets.with_columns([pl.col("true_emission_velocity").cast(pl.Float32)]) + + # cast Int64 to Int32 + dfpackets = dfpackets.with_columns( + [pl.col(col).cast(pl.Int32) for col in dfpackets.columns if dfpackets[col].dtype == pl.Int64] + ) + + return dfpackets + + +def readfile( + packetsfile: Union[Path, str], + packet_type: Optional[str] = None, + escape_type: Optional[Literal["TYPE_RPKT", "TYPE_GAMMA"]] = None, +) -> pd.DataFrame: + """Read a packet file into a Pandas DataFrame.""" + return readfile_pl(packetsfile, packet_type=packet_type, escape_type=escape_type).collect().to_pandas() + + +def readfile_pl( + packetsfile: Union[Path, str], + packet_type: Optional[str] = None, + escape_type: Optional[Literal["TYPE_RPKT", "TYPE_GAMMA"]] = None, +) -> pl.LazyFrame: + """Read a packets file into a Polars LazyFrame from either a parquet file or a text file (and save .parquet).""" + packetsfile = Path(packetsfile) + packetsfileparquet = at.stripallsuffixes(packetsfile).with_suffix(".out.parquet") + packetsfiletext = ( + packetsfile + if packetsfile.suffixes in [[".out"], [".out", ".gz"], [".out", ".xz"], [".out", ".lz4"]] + else at.firstexisting([at.stripallsuffixes(packetsfile).with_suffix(".out")], tryzipped=True) + ) + + write_parquet = True # will be set False if parquet file is read + + dfpackets = None + if packetsfile == packetsfileparquet and os.path.getmtime(packetsfileparquet) > calendar.timegm( + time_lastschemachange + ): + try: + dfpackets = pl.scan_parquet(packetsfileparquet) + write_parquet = False + except Exception as exc: + print(exc) + print(f"Error occured in file {packetsfile}. Reading from text version.") + + if dfpackets is None: + dfpackets = readfile_text(packetsfiletext).lazy() + + if "t_arrive_d" not in dfpackets.columns: + dfpackets = dfpackets.with_columns( + [ + ( + ( + pl.col("escape_time") + - ( + pl.col("posx") * pl.col("dirx") + + pl.col("posy") * pl.col("diry") + + pl.col("posz") * pl.col("dirz") + ) + / 29979245800.0 + ) + / 86400.0 + ).alias("t_arrive_d"), + ] + ) + + if write_parquet: + print(f"Saving {packetsfileparquet}") + dfpackets = dfpackets.sort(by=["type_id", "escape_type_id", "t_arrive_d"]) + dfpackets.collect().write_parquet(packetsfileparquet, compression="zstd", statistics=True) + dfpackets = pl.scan_parquet(packetsfileparquet) + + if escape_type is not None: + assert packet_type is None or packet_type == "TYPE_ESCAPE" + dfpackets = dfpackets.filter( + (pl.col("type_id") == type_ids["TYPE_ESCAPE"]) & (pl.col("escape_type_id") == type_ids[escape_type]) + ) + elif packet_type is not None and packet_type: + dfpackets = dfpackets.filter(pl.col("type_id") == type_ids[packet_type]) + + return dfpackets + + +def get_packetsfilepaths(modelpath: Union[str, Path], maxpacketfiles: Optional[int] = None) -> list[Path]: + nprocs = at.get_nprocs(modelpath) + + searchfolders = [Path(modelpath, "packets"), Path(modelpath)] + # in descending priority (based on speed of reading) + suffix_priority = [".out.parquet", ".out.zst", ".out.lz4", ".out.zst", ".out", ".out.gz", ".out.xz"] + packetsfiles = [] + + for rank in range(nprocs + 1): + name_nosuffix = f"packets00_{rank:04d}" + found_rank = False + for suffix in suffix_priority: + if found_rank: + break + for folderpath in searchfolders: + if (folderpath / name_nosuffix).with_suffix(suffix).is_file(): + packetsfiles.append((folderpath / name_nosuffix).with_suffix(suffix)) + found_rank = True + break + + if found_rank and rank >= nprocs: + print(f"WARNING: nprocs is {nprocs} but file {packetsfiles[-1]} exists") + packetsfiles = packetsfiles[:-1] + elif not found_rank and rank < nprocs: + print(f"WARNING: packets file for rank {rank} was not found.") + + if maxpacketfiles is not None and len(packetsfiles) >= maxpacketfiles: + break + + if maxpacketfiles is not None and nprocs > maxpacketfiles: + print(f"Reading from the first {maxpacketfiles} of {len(packetsfiles)} packets files") + else: + print(f"Reading from {len(packetsfiles)} packets files") + + return packetsfiles + + +def get_packets_pl( + modelpath: Union[str, Path], + maxpacketfiles: Optional[int] = None, + packet_type: Optional[str] = None, + escape_type: Optional[Literal["TYPE_RPKT", "TYPE_GAMMA"]] = None, +) -> tuple[int, pl.LazyFrame]: + if escape_type is not None: + assert packet_type in [None, "TYPE_ESCAPE"] + if packet_type is None: + packet_type = "TYPE_ESCAPE" + + packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles) + + nprocs_read = len(packetsfiles) + allescrpktfile_parquet = Path(modelpath) / "packets_rpkt_escaped.parquet" + + write_allpkts_parquet = False + pldfpackets = None + if maxpacketfiles is None and escape_type == "TYPE_RPKT": + if allescrpktfile_parquet.is_file() and os.path.getmtime(allescrpktfile_parquet) > calendar.timegm( + time_lastschemachange + ): + print(f"Reading from {allescrpktfile_parquet}") + pldfpackets = pl.scan_parquet(allescrpktfile_parquet) + else: + write_allpkts_parquet = True + + if pldfpackets is None: + pldfpackets = pl.concat( + [ + at.packets.readfile_pl(packetsfile, packet_type=packet_type, escape_type=escape_type) + for packetsfile in packetsfiles + ], + how="vertical", + rechunk=False, + ) + + pldfpackets = bin_packet_directions_lazypolars(modelpath, pldfpackets) + + if write_allpkts_parquet: + print(f"Saving {allescrpktfile_parquet}") + # pldfpackets = pldfpackets.sort(by=["type_id", "escape_type_id", "t_arrive_d"]) + pldfpackets.collect(streaming=True).write_parquet( + allescrpktfile_parquet, + compression="zstd", + row_group_size=1024 * 1024, + statistics=True, + ) + + return nprocs_read, pldfpackets + + +def get_directionbin( + dirx: float, diry: float, dirz: float, nphibins: int, ncosthetabins: int, syn_dir: Sequence[float] +) -> int: + dirmag = np.sqrt(dirx**2 + diry**2 + dirz**2) + pkt_dir = [dirx / dirmag, diry / dirmag, dirz / dirmag] + costheta = np.dot(pkt_dir, syn_dir) + thetabin = int((costheta + 1.0) / 2.0 * ncosthetabins) + + xhat = np.array([1.0, 0.0, 0.0]) + vec1 = np.cross(pkt_dir, syn_dir) + vec2 = np.cross(xhat, syn_dir) + cosphi = np.dot(vec1, vec2) / at.vec_len(vec1) / at.vec_len(vec2) + + vec3 = np.cross(vec2, syn_dir) + testphi = np.dot(vec1, vec3) + + phibin = ( + int(math.acos(cosphi) / 2.0 / np.pi * nphibins) + if testphi > 0 + else int((math.acos(cosphi) + np.pi) / 2.0 / np.pi * nphibins) + ) + + return (thetabin * nphibins) + phibin + + +def bin_packet_directions_lazypolars(modelpath: Union[Path, str], dfpackets: pl.LazyFrame) -> pl.LazyFrame: + nphibins = at.get_viewingdirection_phibincount() + ncosthetabins = at.get_viewingdirection_costhetabincount() + + syn_dir = at.get_syn_dir(Path(modelpath)) + xhat = np.array([1.0, 0.0, 0.0]) + vec2 = np.cross(xhat, syn_dir) + + dfpackets = dfpackets.with_columns( + (pl.col("dirx") ** 2 + pl.col("diry") ** 2 + pl.col("dirz") ** 2).sqrt().alias("dirmag"), + ) + dfpackets = dfpackets.with_columns( + ( + (pl.col("dirx") * syn_dir[0] + pl.col("diry") * syn_dir[1] + pl.col("dirz") * syn_dir[2]) / pl.col("dirmag") + ).alias("costheta"), + ) + dfpackets = dfpackets.with_columns( + ((pl.col("costheta") + 1) / 2.0 * ncosthetabins).cast(pl.Int64).alias("costhetabin"), + ) + dfpackets = dfpackets.with_columns( + ((pl.col("diry") * syn_dir[2] - pl.col("dirz") * syn_dir[1]) / pl.col("dirmag")).alias("vec1_x"), + ((pl.col("dirz") * syn_dir[0] - pl.col("dirx") * syn_dir[2]) / pl.col("dirmag")).alias("vec1_y"), + ((pl.col("dirx") * syn_dir[1] - pl.col("diry") * syn_dir[0]) / pl.col("dirmag")).alias("vec1_z"), + ) + + dfpackets = dfpackets.with_columns( + ( + (pl.col("vec1_x") * vec2[0] + pl.col("vec1_y") * vec2[1] + pl.col("vec1_z") * vec2[2]) + / (pl.col("vec1_x") ** 2 + pl.col("vec1_y") ** 2 + pl.col("vec1_z") ** 2).sqrt() + / float(np.linalg.norm(vec2)) + ).alias("cosphi"), + ) + + # vec1 = dir cross syn_dir + dfpackets = dfpackets.with_columns( + ((pl.col("diry") * syn_dir[2] - pl.col("dirz") * syn_dir[1]) / pl.col("dirmag")).alias("vec1_x"), + ((pl.col("dirz") * syn_dir[0] - pl.col("dirx") * syn_dir[2]) / pl.col("dirmag")).alias("vec1_y"), + ((pl.col("dirx") * syn_dir[1] - pl.col("diry") * syn_dir[0]) / pl.col("dirmag")).alias("vec1_z"), + ) + + vec3 = np.cross(vec2, syn_dir) + + # arr_testphi = np.dot(arr_vec1, vec3) + dfpackets = dfpackets.with_columns( + ( + (pl.col("vec1_x") * vec3[0] + pl.col("vec1_y") * vec3[1] + pl.col("vec1_z") * vec3[2]) / pl.col("dirmag") + ).alias("testphi"), + ) + + dfpackets = dfpackets.with_columns( + ( + pl.when(pl.col("testphi") > 0) + .then(pl.col("cosphi").arccos() / 2.0 / np.pi * nphibins) + .otherwise((pl.col("cosphi").arccos() + np.pi) / 2.0 / np.pi * nphibins) + ) + .cast(pl.Int64) + .alias("phibin"), + ) + dfpackets = dfpackets.with_columns( + (pl.col("costhetabin") * nphibins + pl.col("phibin")).alias("dirbin"), + ) + + return dfpackets + + +def bin_packet_directions(modelpath: Union[Path, str], dfpackets: pd.DataFrame) -> pd.DataFrame: + nphibins = at.get_viewingdirection_phibincount() + ncosthetabins = at.get_viewingdirection_costhetabincount() + + syn_dir = at.get_syn_dir(Path(modelpath)) + xhat = np.array([1.0, 0.0, 0.0]) + vec2 = np.cross(xhat, syn_dir) + + pktdirvecs = dfpackets[["dirx", "diry", "dirz"]].to_numpy() + + # normalise. might not be needed + dirmags = np.linalg.norm(pktdirvecs, axis=1) + pktdirvecs /= np.array([dirmags, dirmags, dirmags]).transpose() + + costheta = np.dot(pktdirvecs, syn_dir) + arr_costhetabin = ((costheta + 1) / 2.0 * ncosthetabins).astype(int) + dfpackets["costhetabin"] = arr_costhetabin + + arr_vec1 = np.cross(pktdirvecs, syn_dir) + arr_cosphi = np.dot(arr_vec1, vec2) / np.linalg.norm(arr_vec1, axis=1) / np.linalg.norm(vec2) + vec3 = np.cross(vec2, syn_dir) + arr_testphi = np.dot(arr_vec1, vec3) + + arr_phibin = np.zeros(len(pktdirvecs), dtype=int) + filta = arr_testphi > 0 + arr_phibin[filta] = np.arccos(arr_cosphi[filta]) / 2.0 / np.pi * nphibins + filtb = np.invert(filta) + arr_phibin[filtb] = (np.arccos(arr_cosphi[filtb]) + np.pi) / 2.0 / np.pi * nphibins + dfpackets["phibin"] = arr_phibin + dfpackets["arccoscosphi"] = np.arccos(arr_cosphi) + + dfpackets["dirbin"] = (arr_costhetabin * nphibins) + arr_phibin + + assert np.all(dfpackets["dirbin"] < at.get_viewingdirectionbincount()) + + return dfpackets + + +def make_3d_histogram_from_packets(modelpath, timestep_min, timestep_max=None, em_time=True): + if timestep_max is None: + timestep_max = timestep_min + modeldata, _, vmax_cms = at.inputmodel.get_modeldata_tuple(modelpath) + + timeminarray = at.get_timestep_times_float(modelpath=modelpath, loc="start") + # timedeltaarray = at.get_timestep_times_float(modelpath=modelpath, loc="delta") + timemaxarray = at.get_timestep_times_float(modelpath=modelpath, loc="end") + + # timestep = 63 # 82 73 #63 #54 46 #27 + # print([(ts, time) for ts, time in enumerate(timeminarray)]) + if em_time: + print("Binning by packet emission time") + else: + print("Binning by packet arrival time") + + packetsfiles = at.packets.get_packetsfilepaths(modelpath) + + emission_position3d = [[], [], []] + e_rf = [] + e_cmf = [] + + for packetsfile in packetsfiles: + # for npacketfile in range(0, 1): + dfpackets = at.packets.readfile(packetsfile) + at.packets.add_derived_columns(dfpackets, modelpath, ["emission_velocity"]) + dfpackets = dfpackets.dropna(subset=["emission_velocity"]) # drop rows where emission_vel is NaN + + only_packets_0_scatters = False + if only_packets_0_scatters: + print("Only using packets with 0 scatters") + # print(dfpackets[['scat_count', 'interactions', 'nscatterings']]) + dfpackets = dfpackets.query("nscatterings == 0") + + # print(dfpackets[['emission_velocity', 'em_velx', 'em_vely', 'em_velz']]) + # select only type escape and type r-pkt (don't include gamma-rays) + dfpackets = dfpackets.query( + f'type_id == {type_ids["TYPE_ESCAPE"]} and escape_type_id == {type_ids["TYPE_RPKT"]}' + ) + if em_time: + dfpackets = dfpackets.query("@timeminarray[@timestep_min] < em_time/@DAY < @timemaxarray[@timestep_max]") + else: # packet arrival time + dfpackets = dfpackets.query("@timeminarray[@timestep_min] < t_arrive_d < @timemaxarray[@timestep_max]") + + emission_position3d[0].extend(list(dfpackets["em_velx"] / CLIGHT)) + emission_position3d[1].extend(list(dfpackets["em_vely"] / CLIGHT)) + emission_position3d[2].extend(list(dfpackets["em_velz"] / CLIGHT)) + + e_rf.extend(list(dfpackets["e_rf"])) + e_cmf.extend(list(dfpackets["e_cmf"])) + + emission_position3d = np.array(emission_position3d) + weight_by_energy = True + if weight_by_energy: + e_rf = np.array(e_rf) + e_cmf = np.array(e_cmf) + # weights = e_rf + weights = e_cmf + else: + weights = None + + print(emission_position3d.shape) + print(emission_position3d[0].shape) + + # print(emission_position3d) + grid_3d, _, _, _ = make_3d_grid(modeldata, vmax_cms) + print(grid_3d.shape) + # https://stackoverflow.com/questions/49861468/binning-random-data-to-regular-3d-grid-with-unequal-axis-lengths + hist, _ = np.histogramdd(emission_position3d.T, [np.append(ax, np.inf) for ax in grid_3d], weights=weights) + # print(hist.shape) + if weight_by_energy: + # Divide binned energies by number of processes and by length of timestep + hist = ( + hist / len(packetsfiles) / (timemaxarray[timestep_max] - timeminarray[timestep_min]) + ) # timedeltaarray[timestep] # histogram weighted by energy + # - need to divide by number of processes + # and length of timestep(s) + + # # print histogram coordinates + # coords = np.nonzero(hist) + # for i, j, k in zip(*coords): + # print(f'({grid_3d[0][i]}, {grid_3d[1][j]}, {grid_3d[2][k]}): {hist[i][j][k]}') + + return hist + + +def make_3d_grid(modeldata, vmax_cms): + # modeldata, _, vmax_cms = at.inputmodel.get_modeldata_tuple(modelpath) + grid = round(len(modeldata["inputcellid"]) ** (1.0 / 3.0)) + xgrid = np.zeros(grid) + vmax = vmax_cms / CLIGHT + i = 0 + for _z in range(0, grid): + for _y in range(0, grid): + for x in range(0, grid): + xgrid[x] = -vmax + 2 * x * vmax / grid + i += 1 + + x, y, z = np.meshgrid(xgrid, xgrid, xgrid) + grid_3d = np.array([xgrid, xgrid, xgrid]) + # grid_Te = np.zeros((grid, grid, grid)) + # print(grid_Te.shape) + return grid_3d, x, y, z + + +def get_mean_packet_emission_velocity_per_ts( + modelpath, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT", maxpacketfiles=None, escape_angles=None +) -> pd.DataFrame: + packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles=maxpacketfiles) + nprocs_read = len(packetsfiles) + assert nprocs_read > 0 + + timearray = at.get_timestep_times_float(modelpath=modelpath, loc="mid") + arr_timedelta = at.get_timestep_times_float(modelpath=modelpath, loc="delta") + timearrayplusend = np.concatenate([timearray, [timearray[-1] + arr_timedelta[-1]]]) + + dfpackets_escape_velocity_and_arrive_time = pd.DataFrame + emission_data = pd.DataFrame( + {"t_arrive_d": timearray, "mean_emission_velocity": np.zeros_like(timearray, dtype=float)} + ) + + for i, packetsfile in enumerate(packetsfiles): + dfpackets = at.packets.readfile(packetsfile, packet_type=packet_type, escape_type=escape_type) + at.packets.add_derived_columns(dfpackets, modelpath, ["emission_velocity"]) + if escape_angles is not None: + dfpackets = at.packets.bin_packet_directions(modelpath, dfpackets) + dfpackets = dfpackets.query("dirbin == @escape_angles") + + if i == 0: # make new df + dfpackets_escape_velocity_and_arrive_time = dfpackets[["t_arrive_d", "emission_velocity"]] + else: # append to df + # dfpackets_escape_velocity_and_arrive_time = dfpackets_escape_velocity_and_arrive_time.append( + # other=dfpackets[["t_arrive_d", "emission_velocity"]], ignore_index=True + # ) + dfpackets_escape_velocity_and_arrive_time = pd.concat( + [dfpackets_escape_velocity_and_arrive_time, dfpackets[["t_arrive_d", "emission_velocity"]]], + ignore_index=True, + ) + + print(dfpackets_escape_velocity_and_arrive_time) + binned = pd.cut( + dfpackets_escape_velocity_and_arrive_time["t_arrive_d"], timearrayplusend, labels=False, include_lowest=True + ) + for binindex, emission_velocity in ( + dfpackets_escape_velocity_and_arrive_time.groupby(binned)["emission_velocity"].mean().iteritems() + ): + emission_data["mean_emission_velocity"][binindex] += emission_velocity # / 2.99792458e10 + + return emission_data + + +def bin_and_sum( + df: Union[pl.DataFrame, pl.LazyFrame], + bincol: str, + bins: list[Union[float, int]], + sumcols: Optional[list[str]] = None, + getcounts: bool = False, +) -> pl.DataFrame: + """bins is a list of lower edges, and the final upper edge""" + + # Polars method + df = df.with_columns( + [ + ( + df.select(bincol) + .lazy() + .collect()[bincol] + .cut( + bins=list(bins), + category_label=bincol + "_bin", + maintain_order=True, + ) + .get_column(bincol + "_bin") + .cast(pl.Int32) + - 1 # subtract 1 because the returned index 0 is the bin below the start of the first supplied bin + ) + ] + ) + + if sumcols is not None: + aggs = [pl.col(col).sum().alias(col + "_sum") for col in sumcols] + + if getcounts: + aggs.append(pl.col(bincol).count().alias("count")) + + wlbins = df.groupby(bincol + "_bin").agg(aggs).lazy().collect() + + # now we will include the empty bins + dfout = pl.DataFrame(pl.Series(bincol + "_bin", np.arange(0, len(bins) - 1), dtype=pl.Int32)) + dfout = dfout.join(wlbins, how="left", on=bincol + "_bin").fill_null(0.0) + + # pandas method + + # dfout2 = pd.DataFrame({bincol + "_bin": np.arange(0, len(bins) - 1)}) + # if isinstance(df, pl.DataFrame): + # df2 = df.to_pandas(use_pyarrow_extension_array=True) + + # pdbins = pd.cut( + # x=df2[bincol], + # bins=bins, + # right=True, + # labels=range(len(bins) - 1), + # include_lowest=True, + # ) + + # if sumcols is not None: + # for col in sumcols: + # # dfout = dfout.with_columns( + # # [pl.Series(col + "_sum", df[col].groupby(pdbins).sum().values) for col in sumcols] + # # ) + # dfout2[col + "_sum"] = df2[col].groupby(pdbins).sum().values + # if getcounts: + # # dfout = dfout.with_columns([pl.Series("count", df[bincol].groupby(pdbins).count().values)]) + # dfout2["count"] = df2[bincol].groupby(pdbins).count().values + + return dfout diff --git a/artistools/packets/packetsplots.py b/artistools/packets/packetsplots.py index 5af50ccc3..af2c892f7 100644 --- a/artistools/packets/packetsplots.py +++ b/artistools/packets/packetsplots.py @@ -1,4 +1,3 @@ -#!/usr/bin/env python3 from pathlib import Path import matplotlib.pyplot as plt @@ -35,10 +34,7 @@ def make_2d_packets_plot_imshow(modelpath, timestep_min, timestep_max): time_upper = timemaxarray[timestep_max] title = f"{time_lower:.2f} - {time_upper:.2f} days" print(f"plotting packets between {title}") - if em_time: - escapetitle = "pktemissiontime" - else: - escapetitle = "pktarrivetime" + escapetitle = "pktemissiontime" if em_time else "pktarrivetime" title = title + "\n" + escapetitle plot_axes_list = ["xz", "xy"] @@ -75,7 +71,14 @@ def make_2d_packets_plot_pyvista(modelpath, timestep): mesh["energy [erg/s]"] = hist.ravel(order="F") # print(max(mesh['energy [erg/s]'])) - sargs = dict(height=0.75, vertical=True, position_x=0.04, position_y=0.1, title_font_size=22, label_font_size=25) + sargs = { + "height": 0.75, + "vertical": True, + "position_x": 0.04, + "position_y": 0.1, + "title_font_size": 22, + "label_font_size": 25, + } pv.set_plot_theme("document") # set white background p = pv.Plotter() diff --git a/artistools/plottools.py b/artistools/plottools.py index d25ea0394..9cdc369b1 100644 --- a/artistools/plottools.py +++ b/artistools/plottools.py @@ -1,10 +1,14 @@ +import sys + import matplotlib.pyplot as plt -import matplotlib.ticker as ticker import numpy as np +from matplotlib import ticker + +from .configuration import get_config -import artistools as at -plt.style.use("file://" + str(at.get_config()["path_artistools_dir"] / "matplotlibrc")) +def set_mpl_style() -> None: + plt.style.use("file://" + str(get_config()["path_artistools_dir"] / "matplotlibrc")) class ExponentLabelFormatter(ticker.ScalarFormatter): @@ -138,7 +142,7 @@ def imshow_init_for_artis_grid(ngrid, vmax, plot_variable_3d_array, plot_axes="x plot_axes_choices = ["xy", "xz"] if plot_axes not in plot_axes_choices: print(f"Choose plot axes from {plot_axes_choices}") - quit() + sys.exit(1) for z in range(0, ngrid): for y in range(0, ngrid): @@ -146,9 +150,8 @@ def imshow_init_for_artis_grid(ngrid, vmax, plot_variable_3d_array, plot_axes="x if plot_axes == "xy": if z == round(ngrid / 2) - 1: data[y, x] = plot_variable_3d_array[x, y, z] - elif plot_axes == "xz": - if y == round(ngrid / 2) - 1: - data[z, x] = plot_variable_3d_array[x, y, z] + elif plot_axes == "xz" and y == round(ngrid / 2) - 1: + data[z, x] = plot_variable_3d_array[x, y, z] return data, extent diff --git a/artistools/radfield.py b/artistools/radfield.py old mode 100644 new mode 100755 index 94d7abd3c..ac990c594 --- a/artistools/radfield.py +++ b/artistools/radfield.py @@ -3,7 +3,6 @@ import math import multiprocessing import os -import sys from functools import lru_cache from pathlib import Path from typing import Optional @@ -11,17 +10,9 @@ import matplotlib.pyplot as plt import numpy as np import pandas as pd -from astropy import constants as const from astropy import units as u import artistools as at -import artistools.estimators -import artistools.nltepops -import artistools.spectra - -# import re -# from itertools import chain -# import matplotlib.patches as mpatches H = 6.6260755e-27 # Planck constant [erg s] KB = 1.38064852e-16 # Boltzmann constant [erg/K] @@ -42,7 +33,7 @@ def read_files(modelpath, timestep=-1, modelgridindex=-1): for mpirank in mpiranklist: radfieldfilename = f"radfield_{mpirank:04d}.out" radfieldfilepath = Path(folderpath, radfieldfilename) - radfieldfilepath = at.firstexisting(radfieldfilename, path=folderpath, tryzipped=True) + radfieldfilepath = at.firstexisting(radfieldfilename, folder=folderpath, tryzipped=True) if modelgridindex > -1: filesize = Path(radfieldfilepath).stat().st_size / 1024 / 1024 @@ -52,16 +43,15 @@ def read_files(modelpath, timestep=-1, modelgridindex=-1): # radfielddata_thisfile[['modelgridindex', 'timestep']].apply(pd.to_numeric) if timestep >= 0: - radfielddata_thisfile.query("timestep==@timestep", inplace=True) + radfielddata_thisfile = radfielddata_thisfile.query("timestep==@timestep") if modelgridindex >= 0: - radfielddata_thisfile.query("modelgridindex==@modelgridindex", inplace=True) + radfielddata_thisfile = radfielddata_thisfile.query("modelgridindex==@modelgridindex") if not radfielddata_thisfile.empty: if timestep >= 0 and modelgridindex >= 0: return radfielddata_thisfile - else: - radfielddata = radfielddata.append(radfielddata_thisfile.copy(), ignore_index=True) + radfielddata = radfielddata.append(radfielddata_thisfile.copy(), ignore_index=True) return radfielddata @@ -84,7 +74,7 @@ def select_bin(radfielddata, nu=None, lambda_angstroms=None, modelgridindex=None return dfselected.iloc[0].bin_num, dfselected.iloc[0].nu_lower, dfselected.iloc[0].nu_upper -def get_binaverage_field(axis, radfielddata, modelgridindex=None, timestep=None): +def get_binaverage_field(radfielddata, modelgridindex=None, timestep=None): """Get the dJ/dlambda constant average estimators of each bin.""" # exclude the global fit parameters and detailed lines with negative "bin_num" bindata = radfielddata.copy().query( @@ -93,7 +83,7 @@ def get_binaverage_field(axis, radfielddata, modelgridindex=None, timestep=None) + (" & timestep==@timestep" if timestep else "") ) - arr_lambda = 2.99792458e18 / bindata["nu_upper"].values + arr_lambda = 2.99792458e18 / bindata["nu_upper"].to_numpy() bindata["dlambda"] = bindata.apply(lambda row: 2.99792458e18 * (1 / row["nu_lower"] - 1 / row["nu_upper"]), axis=1) @@ -102,7 +92,7 @@ def get_binaverage_field(axis, radfielddata, modelgridindex=None, timestep=None) row["J"] / row["dlambda"] if (not math.isnan(row["J"] / row["dlambda"]) and row["T_R"] >= 0) else 0.0 ), axis=1, - ).values + ).to_numpy() # add the starting point arr_lambda = np.insert(arr_lambda, 0, 2.99792458e18 / bindata["nu_lower"].iloc[0]) @@ -223,8 +213,8 @@ def plot_line_estimators(axis, radfielddata, xmin, xmax, modelgridindex=None, ti + (" & timestep==@timestep" if timestep else "") )[["nu_upper", "J_nu_avg"]] - radfielddataselected.eval("lambda_angstroms = 2.99792458e18 / nu_upper", inplace=True) - radfielddataselected.eval("Jb_lambda = J_nu_avg * (nu_upper ** 2) / 2.99792458e18", inplace=True) + radfielddataselected = radfielddataselected.eval("lambda_angstroms = 2.99792458e18 / nu_upper") + radfielddataselected = radfielddataselected.eval("Jb_lambda = J_nu_avg * (nu_upper ** 2) / 2.99792458e18") ymax = radfielddataselected["Jb_lambda"].max() @@ -256,7 +246,7 @@ def plot_specout( elif os.path.isfile(specfilename): modelpath = Path(specfilename).parent - dfspectrum = at.spectra.get_spectrum(modelpath, timestep) + dfspectrum = at.spectra.get_spectrum(modelpath=modelpath, timestepmin=modelpath)[-1] label = "Emergent spectrum" if scale_factor is not None: label += " (scaled)" @@ -293,7 +283,7 @@ def sigma_bf(nu): nu_factor = nu / nu_threshold if nu_factor < phixstable[0, 0]: return 0.0 - elif nu_factor > phixstable[-1, 0]: + if nu_factor > phixstable[-1, 0]: # return 0. return phixstable[-1, 1] * math.pow(phixstable[-1, 0] / nu_factor, 3) @@ -358,7 +348,7 @@ def get_recombination_emission( if use_lte_pops: upper_level_popfactor_sum = 0 - for upperlevelnum, upperlevel in lower_ion_data.levels[:200].iterrows(): + for _upperlevelnum, upperlevel in lower_ion_data.levels[:200].iterrows(): upper_level_popfactor_sum += upperlevel.g * math.exp(-upperlevel.energy_ev * EV / KB / T_e) else: dfnltepops = at.nltepops.read_files(modelpath, modelgridindex=modelgridindex, timestep=timestep) @@ -593,7 +583,7 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times # fieldlabel = f'Summed recombination' # axes[2].plot(arraylambda_angstrom_recomb, J_lambda_recomb_total, label=fieldlabel) # fieldlist += [(arraylambda_angstrom_recomb, J_lambda_recomb_total, fieldlabel)] - ymax = max(ymax, max(J_lambda_recomb_total)) + ymax = max(ymax, J_lambda_recomb_total) if args.frompackets: ( @@ -616,12 +606,10 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times ) axes[2].plot(arraylambda_angstrom_em, array_jlambda_emission_total, label="Internal radfield from packets") - fieldlist += list( - [ - (arraylambda_angstrom_em, contribrow.array_flambda_emission, contribrow.linelabel) - for contribrow in contribution_list - ] - ) + fieldlist += [ + (arraylambda_angstrom_em, contribrow.array_flambda_emission, contribrow.linelabel) + for contribrow in contribution_list + ] fieldlist += [(arraylambda_angstrom_em, array_jlambda_emission_total, "Total emission")] # fieldlist += [(arr_lambda_fitted, j_lambda_fitted, 'binned field')] @@ -637,7 +625,7 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times ) if levelnum < 5: axes[0].plot( - arraylambda_angstrom_recomb, arr_sigma_bf, label=r"$\sigma_{bf}$" + f"({ionstr} {level.levelname})" + arraylambda_angstrom_recomb, arr_sigma_bf, label=rf"$\sigma_{{bf}}$({ionstr} {level.levelname})" ) lw = 1.0 @@ -671,7 +659,7 @@ def calculate_photoionrates(axes, modelpath, radfielddata, modelgridindex, times def get_binedges(radfielddata): - return [2.99792458e18 / radfielddata["nu_lower"].iloc[1]] + list(2.99792458e18 / radfielddata["nu_upper"][1:]) + return [2.99792458e18 / radfielddata["nu_lower"].iloc[1], *list(2.99792458e18 / radfielddata["nu_upper"][1:])] def plot_celltimestep(modelpath, timestep, outputfile, xmin, xmax, modelgridindex, args, normalised=False): @@ -711,7 +699,7 @@ def plot_celltimestep(modelpath, timestep, outputfile, xmin, xmax, modelgridinde ymax = max(yvalues) if not args.nobandaverage: - arr_lambda, yvalues = get_binaverage_field(axis, radfielddata, modelgridindex=modelgridindex, timestep=timestep) + arr_lambda, yvalues = get_binaverage_field(radfielddata, modelgridindex=modelgridindex, timestep=timestep) axis.step(arr_lambda, yvalues, where="pre", label="Band-average field", color="green", linewidth=1.5) ymax = max([ymax] + [point[1] for point in zip(arr_lambda, yvalues) if xmin <= point[0] <= xmax]) @@ -731,7 +719,7 @@ def plot_celltimestep(modelpath, timestep, outputfile, xmin, xmax, modelgridinde ymax = args.ymax try: - specfilename = at.firstexisting("spec.out", path=modelpath, tryzipped=True) + specfilename = at.firstexisting("spec.out", folder=modelpath, tryzipped=True) except FileNotFoundError: print("Could not find spec.out") args.nospec = True @@ -788,12 +776,11 @@ def plot_celltimestep(modelpath, timestep, outputfile, xmin, xmax, modelgridinde axis.set_xlabel(r"Wavelength ($\mathrm{{\AA}}$)") axis.set_ylabel(r"J$_\lambda$ [{}erg/s/cm$^2$/$\mathrm{{\AA}}$]") - import matplotlib.ticker as ticker + from matplotlib import ticker axis.xaxis.set_minor_locator(ticker.MultipleLocator(base=500)) axis.set_xlim(left=xmin, right=xmax) axis.set_ylim(bottom=0.0, top=ymax) - import artistools.plottools axis.yaxis.set_major_formatter(at.plottools.ExponentLabelFormatter(axis.get_ylabel(), useMathText=True)) @@ -817,7 +804,9 @@ def plot_bin_fitted_field_evolution(axis, radfielddata, nu_line, modelgridindex, lambda x: j_nu_dbb([nu_line], x.W, x.T_R)[0], axis=1 ) - radfielddataselected.eval("Jb_lambda_at_line = Jb_nu_at_line * (@nu_line ** 2) / 2.99792458e18", inplace=True) + radfielddataselected = radfielddataselected.eval( + "Jb_lambda_at_line = Jb_nu_at_line * (@nu_line ** 2) / 2.99792458e18" + ) lambda_angstroms = 2.99792458e18 / nu_line radfielddataselected.plot( @@ -836,8 +825,8 @@ def plot_global_fitted_field_evolution(axis, radfielddata, nu_line, modelgridind lambda x: j_nu_dbb([nu_line], x.W, x.T_R)[0], axis=1 ) - radfielddataselected.eval( - "J_lambda_fullspec_at_line = J_nu_fullspec_at_line * (@nu_line ** 2) / 2.99792458e18", inplace=True + radfielddataselected = radfielddataselected.eval( + "J_lambda_fullspec_at_line = J_nu_fullspec_at_line * (@nu_line ** 2) / 2.99792458e18" ) lambda_angstroms = 2.99792458e18 / nu_line @@ -861,8 +850,8 @@ def plot_line_estimator_evolution( + (" & timestep <= @timestep_max" if timestep_max else "") )[["timestep", "nu_upper", "J_nu_avg"]] - radfielddataselected.eval("lambda_angstroms = 2.99792458e18 / nu_upper", inplace=True) - radfielddataselected.eval("Jb_lambda = J_nu_avg * (nu_upper ** 2) / 2.99792458e18", inplace=True) + radfielddataselected = radfielddataselected.eval("lambda_angstroms = 2.99792458e18 / nu_upper") + radfielddataselected = radfielddataselected.eval("Jb_lambda = J_nu_avg * (nu_upper ** 2) / 2.99792458e18") axis.plot( radfielddataselected["timestep"], @@ -895,15 +884,17 @@ def plot_timeevolution(modelpath, outputfile, modelgridindex, args): time_days = float(at.get_timestep_time(modelpath, timestep)) dftopestimators = radfielddataselected.query("timestep==@timestep and bin_num < -1").copy() - dftopestimators.eval("lambda_angstroms = 2.99792458e18 / nu_upper", inplace=True) - dftopestimators.eval("Jb_lambda = J_nu_avg * (nu_upper ** 2) / 2.99792458e18", inplace=True) - dftopestimators.sort_values(by="Jb_lambda", ascending=False, inplace=True) + dftopestimators = dftopestimators.eval("lambda_angstroms = 2.99792458e18 / nu_upper") + dftopestimators = dftopestimators.eval("Jb_lambda = J_nu_avg * (nu_upper ** 2) / 2.99792458e18") + dftopestimators = dftopestimators.sort_values(by="Jb_lambda", ascending=False) dftopestimators = dftopestimators.iloc[0:nlinesplotted] print(f"Top estimators at timestep {timestep} t={time_days:.1f}") print(dftopestimators) - for ax, bin_num_estimator, nu_line in zip(axes, dftopestimators.bin_num.values, dftopestimators.nu_upper.values): + for ax, bin_num_estimator, nu_line in zip( + axes, dftopestimators.bin_num.to_numpy(), dftopestimators.nu_upper.to_numpy() + ): lambda_angstroms = 2.99792458e18 / nu_line print(f"Selected line estimator with bin_num {bin_num_estimator}, lambda={lambda_angstroms:.1f}") plot_line_estimator_evolution(ax, radfielddataselected, bin_num_estimator, modelgridindex=modelgridindex) @@ -912,7 +903,7 @@ def plot_timeevolution(modelpath, outputfile, modelgridindex, args): plot_global_fitted_field_evolution(ax, radfielddata, nu_line, modelgridindex=modelgridindex) ax.annotate( - r"$\lambda$=" f"{lambda_angstroms:.1f} Å in cell {modelgridindex:d}\n", + rf"$\lambda$={lambda_angstroms:.1f} Å in cell {modelgridindex:d}\n", xy=(0.02, 0.96), xycoords="axes fraction", horizontalalignment="left", @@ -988,10 +979,13 @@ def main(args=None, argsraw=None, **kwargs): parser.set_defaults(**kwargs) args = parser.parse_args(argsraw) - if args.xaxis == "lambda": - defaultoutputfile = Path("plotradfield_cell{modelgridindex:03d}_ts{timestep:03d}.pdf") - else: - defaultoutputfile = Path("plotradfield_cell{modelgridindex:03d}_evolution.pdf") + at.set_mpl_style() + + defaultoutputfile = ( + Path("plotradfield_cell{modelgridindex:03d}_ts{timestep:03d}.pdf") + if args.xaxis == "lambda" + else Path("plotradfield_cell{modelgridindex:03d}_evolution.pdf") + ) if not args.outputfile: args.outputfile = defaultoutputfile @@ -1006,11 +1000,10 @@ def main(args=None, argsraw=None, **kwargs): if args.velocity >= 0.0: modelgridindexlist = [at.inputmodel.get_mgi_of_velocity_kms(modelpath, args.velocity)] + elif args.modelgridindex is None: + modelgridindexlist = [0] else: - if args.modelgridindex is None: - modelgridindexlist = [0] - else: - modelgridindexlist = at.parse_range_list(args.modelgridindex) + modelgridindexlist = at.parse_range_list(args.modelgridindex) timesteplast = len(at.get_timestep_times_float(modelpath)) if args.timedays: @@ -1054,4 +1047,4 @@ def main(args=None, argsraw=None, **kwargs): if __name__ == "__main__": multiprocessing.freeze_support() - sys.exit(main()) + main() diff --git a/artistools/spectra/__init__.py b/artistools/spectra/__init__.py index 71b2f00f3..79ad50fb3 100644 --- a/artistools/spectra/__init__.py +++ b/artistools/spectra/__init__.py @@ -1,27 +1,24 @@ -#!/usr/bin/env python3 """Artistools - spectra related functions.""" -from artistools.spectra.plotspectra import addargs -from artistools.spectra.plotspectra import main -from artistools.spectra.plotspectra import main as plot -from artistools.spectra.spectra import average_phi_bins -from artistools.spectra.spectra import get_exspec_bins -from artistools.spectra.spectra import get_flux_contributions -from artistools.spectra.spectra import get_flux_contributions_from_packets -from artistools.spectra.spectra import get_line_flux -from artistools.spectra.spectra import get_reference_spectrum -from artistools.spectra.spectra import get_res_spectrum -from artistools.spectra.spectra import get_specpol_data -from artistools.spectra.spectra import get_spectrum -from artistools.spectra.spectra import get_spectrum_at_time -from artistools.spectra.spectra import get_spectrum_from_packets -from artistools.spectra.spectra import get_spectrum_from_packets_worker -from artistools.spectra.spectra import get_vspecpol_spectrum -from artistools.spectra.spectra import make_averaged_vspecfiles -from artistools.spectra.spectra import make_virtual_spectra_summed_file -from artistools.spectra.spectra import print_floers_line_ratio -from artistools.spectra.spectra import print_integrated_flux -from artistools.spectra.spectra import read_spec_res -from artistools.spectra.spectra import sort_and_reduce_flux_contribution_list -from artistools.spectra.spectra import stackspectra -from artistools.spectra.spectra import timeshift_fluxscale_co56law -from artistools.spectra.spectra import write_flambda_spectra +from .__main__ import main +from .plotspectra import addargs +from .plotspectra import main as plot +from .spectra import get_exspec_bins +from .spectra import get_flux_contributions +from .spectra import get_flux_contributions_from_packets +from .spectra import get_from_packets +from .spectra import get_line_flux +from .spectra import get_reference_spectrum +from .spectra import get_specpol_data +from .spectra import get_spectrum +from .spectra import get_spectrum_at_time +from .spectra import get_vspecpol_data +from .spectra import get_vspecpol_spectrum +from .spectra import make_averaged_vspecfiles +from .spectra import make_virtual_spectra_summed_file +from .spectra import print_floers_line_ratio +from .spectra import print_integrated_flux +from .spectra import read_spec_res +from .spectra import sort_and_reduce_flux_contribution_list +from .spectra import stackspectra +from .spectra import timeshift_fluxscale_co56law +from .spectra import write_flambda_spectra diff --git a/artistools/spectra/__main__.py b/artistools/spectra/__main__.py index 16f8b0ba9..bbe6298bd 100755 --- a/artistools/spectra/__main__.py +++ b/artistools/spectra/__main__.py @@ -1,6 +1,10 @@ #!/usr/bin/env python3 -import artistools as at -import artistools.spectra.plotspectra +from .plotspectra import main as plot + + +def main() -> None: + plot() + if __name__ == "__main__": - at.spectra.plotspectra.main() + main() diff --git a/artistools/spectra/plotspectra.py b/artistools/spectra/plotspectra.py index 732511e93..85cf7b4eb 100755 --- a/artistools/spectra/plotspectra.py +++ b/artistools/spectra/plotspectra.py @@ -2,8 +2,8 @@ # PYTHON_ARGCOMPLETE_OK """Artistools - spectra plotting functions.""" import argparse -import math import os +import sys from collections.abc import Collection from pathlib import Path from typing import Any @@ -14,25 +14,22 @@ import argcomplete import matplotlib.patches as mpatches import matplotlib.pyplot as plt -import matplotlib.ticker as ticker import numpy as np import pandas as pd -from astropy import constants as const +from matplotlib import ticker from matplotlib.artist import Artist import artistools as at import artistools.packets -import artistools.radfield -from artistools.spectra.spectra import get_reference_spectrum -from artistools.spectra.spectra import get_res_spectrum -from artistools.spectra.spectra import get_specpol_data -from artistools.spectra.spectra import get_spectrum -from artistools.spectra.spectra import get_spectrum_from_packets -from artistools.spectra.spectra import get_vspecpol_spectrum -from artistools.spectra.spectra import make_averaged_vspecfiles -from artistools.spectra.spectra import make_virtual_spectra_summed_file -from artistools.spectra.spectra import print_integrated_flux -from artistools.spectra.spectra import timeshift_fluxscale_co56law +from .spectra import get_from_packets +from .spectra import get_reference_spectrum +from .spectra import get_specpol_data +from .spectra import get_spectrum +from .spectra import get_vspecpol_spectrum +from .spectra import make_averaged_vspecfiles +from .spectra import make_virtual_spectra_summed_file +from .spectra import print_integrated_flux +from .spectra import timeshift_fluxscale_co56law hatches = ["", "x", "-", "\\", "+", "O", ".", "", "x", "*", "\\", "+", "O", "."] # , @@ -40,7 +37,7 @@ def plot_polarisation(modelpath: Path, args) -> None: angle = args.plotviewingangle[0] stokes_params = get_specpol_data(angle=angle, modelpath=modelpath) - stokes_params[args.stokesparam].eval("lambda_angstroms = 2.99792458e18 / nu", inplace=True) + stokes_params[args.stokesparam] = stokes_params[args.stokesparam].eval("lambda_angstroms = 2.99792458e18 / nu") timearray = stokes_params[args.stokesparam].keys()[1:-1] (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( @@ -51,7 +48,7 @@ def plot_polarisation(modelpath: Path, args) -> None: timeavg = (args.timemin + args.timemax) / 2.0 def match_closest_time(reftime): - return str("{0:.4f}".format(min([float(x) for x in timearray], key=lambda x: abs(x - reftime)))) + return str(f"{min([float(x) for x in timearray], key=lambda x: abs(x - reftime)):.4f}") timeavg = match_closest_time(timeavg) @@ -62,10 +59,11 @@ def match_closest_time(reftime): vpkt_config = at.get_vpkt_config(modelpath) - if args.plotvspecpol: - linelabel = rf"{timeavg} days, cos($\theta$) = {vpkt_config['cos_theta'][angle[0]]}" - else: - linelabel = f"{timeavg} days" + linelabel = ( + f"{timeavg} days, cos($\\theta$) = {vpkt_config['cos_theta'][angle[0]]}" + if args.plotvspecpol + else f"{timeavg} days" + ) if args.binflux: new_lambda_angstroms = [] @@ -147,7 +145,7 @@ def plot_reference_spectrum( print(" metadata: " + ", ".join([f"{k}='{v}'" if hasattr(v, "lower") else f"{k}={v}" for k, v in metadata.items()])) - specdata.query("lambda_angstroms > @xmin and lambda_angstroms < @xmax", inplace=True) + specdata = specdata.query("lambda_angstroms > @xmin and lambda_angstroms < @xmax") print_integrated_flux(specdata["f_lambda"], specdata["lambda_angstroms"], distance_megaparsec=metadata["dist_mpc"]) @@ -155,7 +153,7 @@ def plot_reference_spectrum( # specdata = scipy.signal.resample(specdata, 10000) # specdata = specdata.iloc[::3, :].copy() print(f" downsampling to {len(specdata)} points") - specdata.query("index % 3 == 0", inplace=True) + specdata = specdata.query("index % 3 == 0") # clamp negative values to zero # specdata['f_lambda'] = specdata['f_lambda'].apply(lambda x: max(0, x)) @@ -176,7 +174,7 @@ def plot_reference_spectrum( def plot_filter_functions(axis: plt.Axes) -> None: - filter_names = ["U", "B", "V", "R", "I"] + filter_names = ["U", "B", "V", "I"] colours = ["r", "b", "g", "c", "m"] filterdir = os.path.join(at.get_config()["path_artistools_dir"], "data/filters/") @@ -202,22 +200,34 @@ def plot_artis_spectrum( filterfunc: Optional[Callable[[np.ndarray], np.ndarray]] = None, linelabel: Optional[str] = None, plotpacketcount: bool = False, + directionbins: Optional[list[int]] = None, + average_over_phi: bool = False, + average_over_theta: bool = False, + maxpacketfiles: Optional[int] = None, **plotkwargs, -) -> pd.DataFrame: - """Plot an ARTIS output spectrum. The data plotted are also returned as a DataFrame""" - modelpath_or_file = Path(modelpath) - modelpath = Path(modelpath) if Path(modelpath).is_dir() else Path(modelpath).parent +) -> Optional[pd.DataFrame]: + """Plot an ARTIS output spectrum. The data plotted are also returned as a DataFrame.""" - if not Path(modelpath, "input.txt").exists(): - print(f"Skipping '{modelpath}' (no input.txt found. Not an ARTIS folder?)") - return + modelpath = Path(modelpath) + if Path(modelpath).is_file(): # handle e.g. modelpath = 'modelpath/spec.out' + specfilename = Path(modelpath).parts[-1] + print("WARNING: ignoring filename of {specfilename}") + modelpath = Path(modelpath).parent + + if not modelpath.is_dir(): + print(f"WARNING: Skipping because {modelpath} does not exist") + return None + + if directionbins is None: + directionbins = [-1] if plotpacketcount: from_packets = True - for index, axis in enumerate(axes): + + for axindex, axis in enumerate(axes): if args.multispecplot: (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( - modelpath, timedays_range_str=args.timedayslist[index] + modelpath, timedays_range_str=args.timedayslist[axindex] ) else: (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( @@ -226,59 +236,57 @@ def plot_artis_spectrum( modelname = at.get_model_name(modelpath) if timestepmin == timestepmax == -1: - return + return None assert args.timemin is not None assert args.timemax is not None timeavg = (args.timemin + args.timemax) / 2.0 timedelta = (args.timemax - args.timemin) / 2 if linelabel is None: - if len(modelname) < 70: - linelabel = f"{modelname}" - else: - linelabel = f"...{modelname[-67:]}" + linelabel = f"{modelname}" if len(modelname) < 70 else f"...{modelname[-67:]}" if not args.hidemodeltime and not args.multispecplot: # todo: fix this for multispecplot - use args.showtime for now linelabel += f" +{timeavg:.1f}d" if not args.hidemodeltimerange and not args.multispecplot: - linelabel += r" ($\pm$ " + f"{timedelta:.1f}d)" + linelabel += rf" ($\pm$ {timedelta:.1f}d)" # Luke: disabled below because line label has already been formatted with e.g. timeavg values # formatting for a second time makes it impossible to use curly braces in line labels (needed for LaTeX math) # else: # linelabel = linelabel.format(**locals()) + print( + f"====> '{linelabel}' timesteps {timestepmin} to {timestepmax} ({args.timemin:.3f} to {args.timemax:.3f}d)" + ) + print(f" modelpath {modelpath}") + viewinganglespectra: dict[int, pd.DataFrame] = {} + # have to get the spherical average "bin" if directionbins is None + dbins_get = list(directionbins).copy() + if -1 not in dbins_get: + dbins_get.append(-1) + + supxmin, supxmax = axis.get_xlim() if from_packets: - spectrum = get_spectrum_from_packets( + assert args.plotvspecpol is None + viewinganglespectra = get_from_packets( modelpath, args.timemin, args.timemax, - lambda_min=args.xmin, - lambda_max=args.xmax, + lambda_min=supxmin * 0.9, + lambda_max=supxmax * 1.1, use_escapetime=args.use_escapetime, - maxpacketfiles=args.maxpacketfiles, + maxpacketfiles=maxpacketfiles, delta_lambda=args.deltalambda, useinternalpackets=args.internalpackets, getpacketcount=plotpacketcount, + directionbins=dbins_get, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + fnufilterfunc=filterfunc, ) - - if args.outputfile is None: - statpath = Path() - else: - statpath = Path(args.outputfile).resolve().parent else: - spectrum = get_spectrum(modelpath_or_file, timestepmin, timestepmax, fnufilterfunc=filterfunc) - - angles: Collection[int] = args.plotviewingangle - - if args.plotviewingangle: # read specpol res. - for angle in angles: - viewinganglespectra[angle] = get_res_spectrum( - modelpath, timestepmin, timestepmax, angle=angle, fnufilterfunc=filterfunc, args=args - ) - - elif args.plotvspecpol is not None and os.path.isfile(modelpath / "vpkt.txt"): + if args.plotvspecpol is not None: # noqa: PLR5501 # read virtual packet files (after running plotartisspectrum --makevspecpol) vpkt_config = at.get_vpkt_config(modelpath) if vpkt_config["time_limits_enabled"] and ( @@ -288,99 +296,90 @@ def plot_artis_spectrum( f"Timestep out of range of virtual packets: start time {vpkt_config['initial_time']} days " f"end time {vpkt_config['final_time']} days" ) - quit() - - for angle in angles: - viewinganglespectra[angle] = get_vspecpol_spectrum( - modelpath, timeavg, angle, args, fnufilterfunc=filterfunc - ) + sys.exit(1) - spectrum.query("@args.xmin <= lambda_angstroms and lambda_angstroms <= @args.xmax", inplace=True) + viewinganglespectra = { + dirbin: get_vspecpol_spectrum(modelpath, timeavg, dirbin, args, fnufilterfunc=filterfunc) + for dirbin in dbins_get + if dirbin >= 0 + } + else: + viewinganglespectra = get_spectrum( + modelpath=modelpath, + directionbins=dbins_get, + timestepmin=timestepmin, + timestepmax=timestepmax, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + fnufilterfunc=filterfunc, + ) - print( - f"Plotting '{linelabel}' timesteps {timestepmin} to {timestepmax} " - f"({args.timemin:.3f} to {args.timemax:.3f}d)" + dirbin_definitions = ( + at.get_dirbin_labels( + dirbins=directionbins, + modelpath=modelpath, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + ) + if not args.plotvspecpol + else at.get_vspec_dir_labels(modelpath=modelpath, viewinganglelabelunits=args.viewinganglelabelunits) ) - print(f" modelpath {modelname}") - print_integrated_flux(spectrum["f_lambda"], spectrum["lambda_angstroms"]) - - if scale_to_peak: - spectrum["f_lambda_scaled"] = spectrum["f_lambda"] / spectrum["f_lambda"].max() * scale_to_peak - if args.plotvspecpol is not None: - for angle in args.plotvspecpol: - viewinganglespectra[angle]["f_lambda_scaled"] = ( - viewinganglespectra[angle]["f_lambda"] - / viewinganglespectra[angle]["f_lambda"].max() - * scale_to_peak - ) - ycolumnname = "f_lambda_scaled" - else: - ycolumnname = "f_lambda" + for dirbin in directionbins: + dfspectrum = ( + viewinganglespectra[dirbin] + .query("@supxmin * 0.9 <= lambda_angstroms and lambda_angstroms <= @supxmax * 1.1") + .copy() + ) - if plotpacketcount: - ycolumnname = "packetcount" + if dirbin != -1: + linelabel = dirbin_definitions[dirbin] + print(f" direction {dirbin:4d} {dirbin_definitions[dirbin]}") - supxmin, supxmax = axis.get_xlim() + print_integrated_flux(dfspectrum["f_lambda"], dfspectrum["lambda_angstroms"]) - if (args.plotvspecpol is not None and os.path.isfile(modelpath / "vpkt.txt")) or args.plotviewingangle: - for angle in angles: - if args.binflux: - new_lambda_angstroms = [] - binned_flux = [] + if scale_to_peak: + dfspectrum["f_lambda_scaled"] = dfspectrum["f_lambda"] / dfspectrum["f_lambda"].max() * scale_to_peak + if args.plotvspecpol is not None: + for angle in args.plotvspecpol: + viewinganglespectra[angle]["f_lambda_scaled"] = ( + viewinganglespectra[angle]["f_lambda"] + / viewinganglespectra[angle]["f_lambda"].max() + * scale_to_peak + ) - wavelengths = viewinganglespectra[angle]["lambda_angstroms"] - fluxes = viewinganglespectra[angle][ycolumnname] - nbins = 5 + ycolumnname = "f_lambda_scaled" + else: + ycolumnname = "f_lambda" - for i in np.arange(0, len(wavelengths - nbins), nbins): - new_lambda_angstroms.append(wavelengths[i + int(nbins / 2)]) - sum_flux = 0 - for j in range(i, i + nbins): - sum_flux += fluxes[j] - binned_flux.append(sum_flux / nbins) + if plotpacketcount: + ycolumnname = "packetcount" - plt.plot(new_lambda_angstroms, binned_flux) - else: - if args.plotvspecpol: - if args.viewinganglelabelunits == "deg": - viewing_angle = round(math.degrees(math.acos(vpkt_config["cos_theta"][angle]))) - linelabel = rf"$\theta$ = {viewing_angle}$^\circ$" if index == 0 else None - elif args.viewinganglelabelunits == "rad": - linelabel = rf"cos($\theta$) = {vpkt_config['cos_theta'][angle]}" if index == 0 else None - else: - linelabel = f"bin number {angle}" - if args.average_over_phi_angle: - ( - costheta_viewing_angle_bins, - phi_viewing_angle_bins, - ) = at.get_viewinganglebin_definitions() - assert angle % at.get_viewingdirection_phibincount() == 0 - linelabel = costheta_viewing_angle_bins[int(angle // 10)] - elif args.average_over_theta_angle: - ( - costheta_viewing_angle_bins, - phi_viewing_angle_bins, - ) = at.get_viewinganglebin_definitions() - assert angle < at.get_viewingdirection_costhetabincount() - linelabel = phi_viewing_angle_bins[int(angle)] - - viewinganglespectra[angle].query( - "@supxmin <= lambda_angstroms and lambda_angstroms <= @supxmax" - ).plot( - x="lambda_angstroms", y=ycolumnname, ax=axis, legend=None, label=linelabel - ) # {timeavg:.2f} days {at.get_model_name(modelpath)} - else: - spectrum.query("@supxmin <= lambda_angstroms and lambda_angstroms <= @supxmax").plot( - x="lambda_angstroms", - y=ycolumnname, - ax=axis, - legend=None, - label=linelabel if index == 0 else None, + if args.binflux: + new_lambda_angstroms = [] + binned_flux = [] + + wavelengths = dfspectrum["lambda_angstroms"] + fluxes = dfspectrum[ycolumnname] + nbins = 5 + + for i in np.arange(0, len(wavelengths - nbins), nbins): + new_lambda_angstroms.append(wavelengths[i + int(nbins / 2)]) + sum_flux = 0 + for j in range(i, i + nbins): + sum_flux += fluxes[j] + binned_flux.append(sum_flux / nbins) + + dfspectrum = pd.DataFrame({"lambda_angstroms": new_lambda_angstroms, ycolumnname: binned_flux}) + + axis.plot( + dfspectrum["lambda_angstroms"], + dfspectrum[ycolumnname], + label=linelabel if axindex == 0 else None, **plotkwargs, ) - return spectrum[["lambda_angstroms", "f_lambda"]] + return dfspectrum[["lambda_angstroms", "f_lambda"]] def make_spectrum_plot( @@ -402,7 +401,8 @@ def make_spectrum_plot( plotkwargs["alpha"] = 0.95 plotkwargs["linestyle"] = args.linestyle[seriesindex] - plotkwargs["color"] = args.color[seriesindex] + if not args.plotviewingangle and not args.plotvspecpol: + plotkwargs["color"] = args.color[seriesindex] if args.dashes[seriesindex]: plotkwargs["dashes"] = args.dashes[seriesindex] if args.linewidth[seriesindex]: @@ -410,36 +410,7 @@ def make_spectrum_plot( seriesdata = pd.DataFrame() - if specpath.is_dir() or specpath.name.endswith(".out"): - # ARTIS model spectrum - # plotkwargs['dash_capstyle'] = dash_capstyleList[artisindex] - if "linewidth" not in plotkwargs: - plotkwargs["linewidth"] = 1.3 - - plotkwargs["linelabel"] = args.label[seriesindex] - - seriesdata = plot_artis_spectrum( - axes, - specpath, - args=args, - scale_to_peak=scale_to_peak, - from_packets=args.frompackets, - filterfunc=filterfunc, - plotpacketcount=args.plotpacketcount, - **plotkwargs, - ) - seriesname = at.get_model_name(specpath) - artisindex += 1 - - elif not specpath.exists() and specpath.parts[0] == "codecomparison": - # timeavg = (args.timemin + args.timemax) / 2. - (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( - specpath, args.timestep, args.timemin, args.timemax, args.timedays - ) - timeavg = args.timedays - artistools.codecomparison.plot_spectrum(specpath, timedays=timeavg, ax=axes[0], **plotkwargs) - refspecindex += 1 - else: + if not Path(specpath).is_dir() and not Path(specpath).exists() and "." in str(specpath): # reference spectrum if "linewidth" not in plotkwargs: plotkwargs["linewidth"] = 1.1 @@ -458,7 +429,7 @@ def make_spectrum_plot( **plotkwargs, ) else: - for _, axis in enumerate(axes): + for axis in axes: supxmin, supxmax = axis.get_xlim() plot_reference_spectrum( specpath, @@ -471,17 +442,54 @@ def make_spectrum_plot( **plotkwargs, ) refspecindex += 1 + elif not specpath.exists() and specpath.parts[0] == "codecomparison": + # timeavg = (args.timemin + args.timemax) / 2. + (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( + specpath, args.timestep, args.timemin, args.timemax, args.timedays + ) + timeavg = args.timedays + artistools.codecomparison.plot_spectrum(specpath, timedays=timeavg, ax=axes[0], **plotkwargs) + refspecindex += 1 + else: + # ARTIS model spectrum + # plotkwargs['dash_capstyle'] = dash_capstyleList[artisindex] + if "linewidth" not in plotkwargs: + plotkwargs["linewidth"] = 1.3 + + plotkwargs["linelabel"] = args.label[seriesindex] + + seriesdata = plot_artis_spectrum( + axes, + specpath, + args=args, + scale_to_peak=scale_to_peak, + from_packets=args.frompackets, + maxpacketfiles=args.maxpacketfiles, + filterfunc=filterfunc, + plotpacketcount=args.plotpacketcount, + directionbins=args.plotviewingangle if not args.plotvspecpol else args.plotvspecpol, + average_over_phi=args.average_over_phi_angle, + average_over_theta=args.average_over_theta_angle, + **plotkwargs, + ) + if seriesdata is not None: + seriesname = at.get_model_name(specpath) + artisindex += 1 if args.write_data and not seriesdata.empty: if dfalldata.empty: - dfalldata = pd.DataFrame(index=seriesdata["lambda_angstroms"].values) + dfalldata = pd.DataFrame(index=seriesdata["lambda_angstroms"].to_numpy()) dfalldata.index.name = "lambda_angstroms" else: - assert np.allclose(dfalldata.index.values, seriesdata["lambda_angstroms"].values) - dfalldata[f"f_lambda.{seriesname}"] = seriesdata["f_lambda"].values + # make sure we can share the same set of wavelengths for this series + assert np.allclose(dfalldata.index.values, seriesdata["lambda_angstroms"].to_numpy()) + dfalldata[f"f_lambda.{seriesname}"] = seriesdata["f_lambda"].to_numpy() seriesindex += 1 + plottedsomething = artisindex > 0 or refspecindex > 0 + assert plottedsomething + for axis in axes: if args.showfilterfunctions: if not args.normalised: @@ -527,16 +535,16 @@ def make_emissionabsorption_plot( args=None, scale_to_peak: Optional[float] = None, ) -> tuple[list[Artist], list[str], Optional[pd.DataFrame]]: - """Plot the emission and absorption by ion for an ARTIS model.""" - print(modelpath) - arraynu = at.misc.get_nu_grid(modelpath) + """Plot the emission and absorption contribution spectra, grouped by ion/line/term for an ARTIS model.""" + modelname = at.get_model_name(modelpath) + + print(f"====> {modelname}") + arraynu = at.get_nu_grid(modelpath) (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( modelpath, args.timestep, args.timemin, args.timemax, args.timedays ) - modelname = at.get_model_name(modelpath) - if timestepmin == timestepmax == -1: print(f"Can't plot {modelname}...skipping") return [], [], None @@ -670,7 +678,7 @@ def make_emissionabsorption_plot( if not args.showemission: plotobjects.extend(absstackplot) - plotobjectlabels.extend(list([x.linelabel for x in contributions_sorted_reduced])) + plotobjectlabels.extend([x.linelabel for x in contributions_sorted_reduced]) # print(plotobjectlabels) # print(len(plotobjectlabels), len(plotobjects)) @@ -702,17 +710,13 @@ def make_emissionabsorption_plot( plotlabel = f"{modelname}\n{args.timemin:.2f}d to {args.timemax:.2f}d" if args.plotviewingangle: - angle = args.plotviewingangle[0] - ( - costheta_viewing_angle_bins, - phi_viewing_angle_bins, - ) = at.get_viewinganglebin_definitions() - - if args.average_over_phi_angle: - assert angle % 10 == 0 - plotlabel += ", " + costheta_viewing_angle_bins[int(angle // 10)] - else: - plotlabel += f", directionbin {args.plotviewingangle[0]}" + dirbin_definitions = at.get_dirbin_labels( + dirbins=args.plotviewingangle, + modelpath=modelpath, + average_over_phi=args.average_over_phi_angle, + average_over_theta=args.average_over_theta_angle, + ) + plotlabel += f", directionbin {dirbin_definitions[args.plotviewingangle[0]]}" if not args.notitle: axis.set_title(plotlabel, fontsize=11) @@ -740,8 +744,6 @@ def make_emissionabsorption_plot( def make_contrib_plot(axes: plt.Axes, modelpath: Path, densityplotyvars: list[str], args) -> None: - import artistools.packets - (timestepmin, timestepmax, args.timemin, args.timemax) = at.get_time_range( modelpath, args.timestep, args.timemin, args.timemax, args.timedays ) @@ -749,11 +751,9 @@ def make_contrib_plot(axes: plt.Axes, modelpath: Path, densityplotyvars: list[st modeldata, _ = at.inputmodel.get_modeldata(modelpath) if args.classicartis: - import artistools.estimators.estimators_classic - modeldata, _ = at.inputmodel.get_modeldata(modelpath) estimators = artistools.estimators.estimators_classic.read_classic_estimators(modelpath, modeldata) - allnonemptymgilist = [modelgridindex for modelgridindex in modeldata.index] + allnonemptymgilist = list(modeldata.index) else: estimators = at.estimators.read_estimators(modelpath=modelpath) @@ -764,19 +764,17 @@ def make_contrib_plot(axes: plt.Axes, modelpath: Path, densityplotyvars: list[st packetsfiles = at.packets.get_packetsfilepaths(modelpath, args.maxpacketfiles) assert args.timemin is not None assert args.timemax is not None - tdays_min = float(args.timemin) - tdays_max = float(args.timemax) + # tdays_min = float(args.timemin) + # tdays_max = float(args.timemax) c_ang_s = 2.99792458e18 - nu_min = c_ang_s / args.xmax - nu_max = c_ang_s / args.xmin - - querystr = "" + # nu_min = c_ang_s / args.xmax + # nu_max = c_ang_s / args.xmin list_lambda: dict[str, list[float]] = {} lists_y: dict[str, list[float]] = {} - for index, packetsfile in enumerate(packetsfiles): - dfpackets = at.packets.readfile(packetsfile, type="TYPE_ESCAPE", escape_type="TYPE_RPKT") + for packetsfile in packetsfiles: + dfpackets = at.packets.readfile(packetsfile, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT") dfpackets_selected = dfpackets.query( "@nu_min <= nu_rf < @nu_max and t_arrive_d >= @tdays_min and t_arrive_d <= @tdays_max", inplace=False @@ -841,10 +839,7 @@ def make_plot(args) -> None: # densityplotyvars = ['emission_velocity', 'Te', 'nne'] # densityplotyvars = ['true_emission_velocity', 'emission_velocity', 'Te', 'nne'] - if args.multispecplot: - nrows = len(args.timedayslist) - else: - nrows = 1 + len(densityplotyvars) + nrows = len(args.timedayslist) if args.multispecplot else 1 + len(densityplotyvars) fig, axes = plt.subplots( nrows=nrows, @@ -878,7 +873,7 @@ def make_plot(args) -> None: else: axes[-1].set_ylabel(r"F$_\lambda$ at 1 Mpc [{}erg/s/cm$^2$/$\mathrm{{\AA}}$]") - for index, axis in enumerate(axes): + for axis in axes: if args.logscale: axis.set_yscale("log") axis.set_xlim(left=args.xmin, right=args.xmax) @@ -929,14 +924,14 @@ def make_plot(args) -> None: plotobjectlabels, loc="upper right", frameon=False, - handlelength=1, + handlelength=2, ncol=legendncol, numpoints=1, fontsize=fs, ) leg.set_zorder(200) - for artist, text in zip(leg.legendHandles, leg.get_texts()): + for artist, text in zip(leg.legend_handles, leg.get_texts()): if hasattr(artist, "get_color"): col = artist.get_color() artist.set_linewidth(2.0) @@ -990,6 +985,7 @@ def make_plot(args) -> None: filenameout = str(args.outputfile).format( time_days_min=args.timemin, time_days_max=args.timemax, directionbins=strdirectionbins ) + # plt.text(6000, (args.ymax * 0.9), f'{round(args.timemin) + 1} days', fontsize='large') if args.showtime and not args.multispecplot: @@ -1085,7 +1081,7 @@ def addargs(parser) -> None: parser.add_argument( "-filtersavgol", nargs=2, - help="Savitzky–Golay filter. Specify the window_length and poly_order.e.g. -filtersavgol 5 3", + help="Savitzky-Golay filter. Specify the window_length and poly_order.e.g. -filtersavgol 5 3", ) parser.add_argument("-timestep", "-ts", dest="timestep", nargs="?", help="First timestep or a range e.g. 45-65") @@ -1284,6 +1280,12 @@ def main(args=None, argsraw=None, **kwargs) -> None: print("WARNING: --average_every_tenth_viewing_angle is deprecated. use --average_over_phi_angle instead") args.average_over_phi_angle = True + at.set_mpl_style() + + assert ( + not args.plotvspecpol or not args.plotviewingangle + ) # choose either virtual packet directions or real packet direction bins + if not args.specpath: args.specpath = [Path(".")] elif isinstance(args.specpath, (str, Path)): # or not not isinstance(args.specpath, Iterable) @@ -1328,7 +1330,7 @@ def main(args=None, argsraw=None, **kwargs) -> None: plot_polarisation(args.specpath[0], args) return - elif args.output_spectra: + if args.output_spectra: for modelpath in args.specpath: at.spectra.write_flambda_spectra(modelpath, args) diff --git a/artistools/spectra/sampleblackbodyfrompacketTR.py b/artistools/spectra/sampleblackbodyfrompacketTR.py index 2402ff95d..b385254e5 100644 --- a/artistools/spectra/sampleblackbodyfrompacketTR.py +++ b/artistools/spectra/sampleblackbodyfrompacketTR.py @@ -42,7 +42,7 @@ 20: "TYPE_NTLEPTON", 32: "TYPE_ESCAPE", } -type_ids = dict((v, k) for k, v in types.items()) +type_ids = {v: k for k, v in types.items()} def sample_planck(T, nu_max_r, nu_min_r): @@ -96,8 +96,8 @@ def planck(nu, T): # nprocs = 100 for npacketfile in range(0, nprocs): dfpackets = at.packets.readfile(packetsfiles[npacketfile]) # , type='TYPE_ESCAPE', escape_type='TYPE_RPKT') - dfpackets = at.packets.get_escaping_packet_angle_bin(modelpath, dfpackets) - dfpackets.query(f'type_id == {type_ids["TYPE_ESCAPE"]} and escape_type_id == {type_ids["TYPE_RPKT"]}', inplace=True) + dfpackets = at.packets.bin_packet_directions(modelpath, dfpackets) + dfpackets = dfpackets.query(f'type_id == {type_ids["TYPE_ESCAPE"]} and escape_type_id == {type_ids["TYPE_RPKT"]}') # print(max(dfpackets['t_arrive_d'])) # print(dfpackets) @@ -123,9 +123,9 @@ def planck(nu, T): # print('initial df:') # print((dfpackets[['escape_time', 'escape_time_d', 't_arrive_d', 'em_TR']])) # print('\n\n\n') - # # quit() + # # sys.exit(1) - for df_index, row in dfpackets_timestep.iterrows(): + for _df_index, row in dfpackets_timestep.iterrows(): TR = row["em_TR"] # if TR not in [100, 140000]: diff --git a/artistools/spectra/spectra.py b/artistools/spectra/spectra.py index 6e3942e4c..2ecb24f89 100644 --- a/artistools/spectra/spectra.py +++ b/artistools/spectra/spectra.py @@ -1,15 +1,12 @@ -#!/usr/bin/env python3 """Artistools - spectra related functions.""" import argparse import math -import multiprocessing import os import re from collections import namedtuple from collections.abc import Collection from collections.abc import Sequence from functools import lru_cache -from functools import partial from pathlib import Path from typing import Any from typing import Callable @@ -18,15 +15,14 @@ from typing import Union import matplotlib as mpl -import matplotlib.pyplot as plt # needed to get the color map +import matplotlib.pyplot as plt import numpy as np import pandas as pd +import polars as pl from astropy import constants as const from astropy import units as u import artistools as at -import artistools.packets -import artistools.radfield fluxcontributiontuple = namedtuple( "fluxcontributiontuple", "fluxcontrib linelabel array_flambda_emission array_flambda_absorption color" @@ -38,15 +34,19 @@ def timeshift_fluxscale_co56law(scaletoreftime: Optional[float], spectime: float # Co56 decay flux scaling assert spectime > 150 return math.exp(float(spectime) / 113.7) / math.exp(scaletoreftime / 113.7) - else: - return 1.0 + + return 1.0 def get_exspec_bins() -> tuple[np.ndarray, np.ndarray, np.ndarray]: MNUBINS = 1000 NU_MIN_R = 1e13 - NU_MAX_R = 5e15 - print("WARNING: assuming {MNUBINS=} {NU_MIN_R=} {NU_MAX_R=}. Check that artis code matches.") + NU_MAX_R = 5e16 + + print( + f" assuming {MNUBINS=} {NU_MIN_R=:.1e} {NU_MAX_R=:.1e}. Check artisoptions.h if you want to exactly match" + " exspec binning." + ) c_ang_s = 2.99792458e18 @@ -80,142 +80,39 @@ def stackspectra( return stackedspectrum -@lru_cache(maxsize=16) -def get_specdata(modelpath: Path, stokesparam: Optional[Literal["I", "Q", "U"]] = None) -> pd.DataFrame: - polarisationdata = False - if Path(modelpath, "specpol.out").is_file(): - specfilename = Path(modelpath) / "specpol.out" - polarisationdata = True - elif Path(modelpath, "specpol.out.xz").is_file(): - specfilename = Path(modelpath) / "specpol.out.xz" - polarisationdata = True - elif Path(modelpath).is_dir(): - specfilename = at.firstexisting("spec.out", path=modelpath, tryzipped=True) - else: - specfilename = modelpath - - if polarisationdata: - # angle = args.plotviewingangle[0] - stokes_params = get_specpol_data(angle=None, modelpath=modelpath) - if stokesparam is not None: - specdata = stokes_params[stokesparam] - else: - specdata = stokes_params["I"] - else: - assert stokesparam is None - print(f"Reading {specfilename}") - specdata = pd.read_csv(specfilename, delim_whitespace=True) - specdata.rename(columns={"0": "nu"}, inplace=True) - - return specdata - - -def get_spectrum( - modelpath: Path, - timestepmin: int, - timestepmax: int = -1, - fnufilterfunc: Optional[Callable[[np.ndarray], np.ndarray]] = None, - modelnumber: Optional[int] = None, -) -> pd.DataFrame: - """Return a pandas DataFrame containing an ARTIS emergent spectrum.""" - if timestepmax < 0: - timestepmax = timestepmin - - specdata = get_specdata(modelpath) - - nu = specdata.loc[:, "nu"].values - arr_tdelta = at.get_timestep_times_float(modelpath, loc="delta") - - f_nu = stackspectra( - [ - (specdata[specdata.columns[timestep + 1]], arr_tdelta[timestep]) - for timestep in range(timestepmin, timestepmax + 1) - ] - ) - - # best to use the filter on this list because it - # has regular sampling - if fnufilterfunc: - print("Applying filter to ARTIS spectrum") - f_nu = fnufilterfunc(f_nu) - - dfspectrum = pd.DataFrame({"nu": nu, "f_nu": f_nu}) - dfspectrum.sort_values(by="nu", ascending=False, inplace=True) - - dfspectrum.eval("lambda_angstroms = @c / nu", local_dict={"c": 2.99792458e18}, inplace=True) - dfspectrum.eval("f_lambda = f_nu * nu / lambda_angstroms", inplace=True) - - # if 'redshifttoz' in args and args.redshifttoz[modelnumber] != 0: - # # plt.plot(dfspectrum['lambda_angstroms'], dfspectrum['f_lambda'], color='k') - # z = args.redshifttoz[modelnumber] - # dfspectrum['lambda_angstroms'] *= (1 + z) - # # plt.plot(dfspectrum['lambda_angstroms'], dfspectrum['f_lambda'], color='r') - # # plt.show() - # # quit() - - return dfspectrum - - def get_spectrum_at_time( modelpath: Path, timestep: int, time: float, args: Optional[argparse.Namespace], - angle: Optional[int] = None, - res_specdata: Optional[dict[int, pd.DataFrame]] = None, - modelnumber: Optional[int] = None, + dirbin: int = -1, + average_over_phi: Optional[bool] = None, + average_over_theta: Optional[bool] = None, ) -> pd.DataFrame: - if angle is not None and angle >= 0: + if dirbin >= 0: if args is not None and args.plotvspecpol and os.path.isfile(modelpath / "vpkt.txt"): - spectrum = get_vspecpol_spectrum(modelpath, time, angle, args) - else: - spectrum = get_res_spectrum(modelpath, timestep, timestep, angle=angle, res_specdata=res_specdata) - else: - spectrum = get_spectrum(modelpath, timestep, timestep, modelnumber=modelnumber) - - return spectrum - - -def get_spectrum_from_packets_worker( - querystr: str, - qlocals: dict[str, Any], - array_lambda: Sequence[float], - array_lambdabinedges: Sequence[float], - packetsfile: Path, - use_escapetime: bool = False, - getpacketcount: bool = False, - betafactor: Optional[float] = None, -) -> tuple[np.ndarray, np.ndarray]: - dfpackets = at.packets.readfile(packetsfile, type="TYPE_ESCAPE", escape_type="TYPE_RPKT").query( - querystr, inplace=False, local_dict=qlocals - ) - - print(f" {packetsfile}: {len(dfpackets)} escaped r-packets matching frequency and arrival time ranges ") + spectrum = get_vspecpol_spectrum(modelpath, time, dirbin, args) + return spectrum - dfpackets.eval("lambda_rf = 2.99792458e18 / nu_rf", inplace=True, local_dict=qlocals) - wl_bins = pd.cut( - x=dfpackets["lambda_rf"], - bins=array_lambdabinedges, - right=True, - labels=range(len(array_lambda)), - include_lowest=True, - ) - - if use_escapetime: - assert betafactor is not None - array_energysum_onefile = dfpackets.e_cmf.groupby(wl_bins).sum().values / betafactor + assert average_over_phi is not None + assert average_over_theta is not None else: - array_energysum_onefile = dfpackets.e_rf.groupby(wl_bins).sum().values + average_over_phi = False + average_over_theta = False + + spectrum = get_spectrum( + modelpath=modelpath, + directionbins=[dirbin], + timestepmin=timestep, + timestepmax=timestep, + average_over_phi=average_over_phi, + average_over_theta=average_over_theta, + )[dirbin] - if getpacketcount: - array_pktcount_onefile = dfpackets.lambda_rf.groupby(wl_bins).count().values - else: - array_pktcount_onefile = None - - return array_energysum_onefile, array_pktcount_onefile + return spectrum -def get_spectrum_from_packets( +def get_from_packets( modelpath: Path, timelowdays: float, timehighdays: float, @@ -226,17 +123,22 @@ def get_spectrum_from_packets( maxpacketfiles: Optional[int] = None, useinternalpackets: bool = False, getpacketcount: bool = False, + directionbins: Optional[Collection[int]] = None, + average_over_phi: bool = False, + average_over_theta: bool = False, + fnufilterfunc: Optional[Callable[[np.ndarray], np.ndarray]] = None, ) -> pd.DataFrame: """Get a spectrum dataframe using the packets files as input.""" assert not useinternalpackets - packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles) + if directionbins is None: + directionbins = [-1] if use_escapetime: - modeldata, _ = at.inputmodel.get_modeldata(Path(packetsfiles[0]).parent) - vmax = modeldata.iloc[-1].velocity_outer * u.km / u.s - betafactor = math.sqrt(1 - (vmax / const.c).decompose().value ** 2) + modeldata, _ = at.inputmodel.get_modeldata(modelpath) + vmax_beta = modeldata.iloc[-1].velocity_outer * 299792.458 + escapesurfacegamma = math.sqrt(1 - vmax_beta**2) else: - betafactor = None + escapesurfacegamma = None nu_min = 2.99792458e18 / lambda_max nu_max = 2.99792458e18 / lambda_min @@ -248,221 +150,264 @@ def get_spectrum_from_packets( else: array_lambdabinedges, array_lambda, delta_lambda = get_exspec_bins() - array_energysum = np.zeros_like(array_lambda, dtype=float) # total packet energy sum of each bin - if getpacketcount: - array_pktcount = np.zeros_like(array_lambda, dtype=int) # number of packets in each bin - timelow = timelowdays * 86400.0 timehigh = timehighdays * 86400.0 - nprocs_read = len(packetsfiles) - querystr = "@nu_min <= nu_rf < @nu_max and trueemissiontype >= 0 and " - if not use_escapetime: - querystr += "@timelow < (escape_time - (posx * dirx + posy * diry + posz * dirz) / 29979245800) < @timehigh" - else: - querystr += "@timelow < (escape_time * @betafactor) < @timehigh" - - processfile = partial( - get_spectrum_from_packets_worker, - querystr, - dict( - nu_min=nu_min, - nu_max=nu_max, - timelow=timelow, - timehigh=timehigh, - betafactor=betafactor, - ), - array_lambda, - array_lambdabinedges, - use_escapetime=use_escapetime, - getpacketcount=getpacketcount, - betafactor=betafactor, + nphibins = at.get_viewingdirection_phibincount() + ncosthetabins = at.get_viewingdirection_costhetabincount() + ndirbins = at.get_viewingdirectionbincount() + + nprocs_read, dfpackets = at.packets.get_packets_pl( + modelpath, maxpacketfiles=maxpacketfiles, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT" ) - if at.get_config()["num_processes"] > 1: - with multiprocessing.Pool(processes=at.get_config()["num_processes"]) as pool: - results = pool.map(processfile, packetsfiles) - pool.close() - pool.join() - pool.terminate() + + if not use_escapetime: + dfpackets = dfpackets.filter( + (float(timelowdays) <= pl.col("t_arrive_d")) & (pl.col("t_arrive_d") <= float(timehighdays)) + ) else: - results = [processfile(p) for p in packetsfiles] + dfpackets = dfpackets.filter( + (timelow <= (pl.col("escape_time") * escapesurfacegamma)) + & ((pl.col("escape_time") * escapesurfacegamma) <= timehigh) + ) + dfpackets = dfpackets.filter((float(nu_min) <= pl.col("nu_rf")) & (pl.col("nu_rf") <= float(nu_max))) - array_energysum = np.ufunc.reduce(np.add, [r[0] for r in results]) - if getpacketcount: - array_pktcount += np.ufunc.reduce(np.add, [r[1] for r in results]) + if fnufilterfunc: + print("Applying filter to ARTIS spectrum") - array_flambda = ( - array_energysum / delta_lambda / (timehigh - timelow) / 4 / math.pi / (u.megaparsec.to("cm") ** 2) / nprocs_read - ) + encol = "e_cmf" if use_escapetime else "e_rf" + getcols = ["nu_rf", encol] + if directionbins != [-1]: + if average_over_phi: + getcols.append("costhetabin") + elif average_over_theta: + getcols.append("phibin") + else: + getcols.append("dirbin") + dfpackets = dfpackets.select(getcols).collect().lazy() + + dfdict = {} + for dirbin in directionbins: + if dirbin == -1: + solidanglefactor = 1.0 + pldfpackets_dirbin_lazy = dfpackets + elif average_over_phi: + assert not average_over_theta + solidanglefactor = ncosthetabins + pldfpackets_dirbin_lazy = dfpackets.filter(pl.col("costhetabin") * 10 == dirbin) + elif average_over_theta: + solidanglefactor = nphibins + pldfpackets_dirbin_lazy = dfpackets.filter(pl.col("phibin") == dirbin) + else: + solidanglefactor = ndirbins + pldfpackets_dirbin_lazy = dfpackets.filter(pl.col("dirbin") == dirbin) + + pldfpackets_dirbin = pldfpackets_dirbin_lazy.with_columns( + [(2.99792458e18 / pl.col("nu_rf")).alias("lambda_angstroms")] + ).select(["lambda_angstroms", encol]) + + dfbinned = at.packets.bin_and_sum( + pldfpackets_dirbin, + bincol="lambda_angstroms", + bins=list(array_lambdabinedges), + sumcols=[encol], + getcounts=getpacketcount, + ) + array_flambda = ( + dfbinned[encol + "_sum"] + / delta_lambda + / (timehigh - timelow) + / (4 * math.pi) + * solidanglefactor + / (u.megaparsec.to("cm") ** 2) + / nprocs_read + ) - dfdict = { - "lambda_angstroms": array_lambda, - "f_lambda": array_flambda, - "energy_sum": array_energysum, - } + if use_escapetime: + assert escapesurfacegamma is not None + array_flambda /= escapesurfacegamma + + if fnufilterfunc: + arr_nu = 2.99792458e18 / array_lambda + array_f_nu = array_flambda * array_lambda / arr_nu + array_f_nu = fnufilterfunc(array_f_nu) + array_flambda = array_f_nu * arr_nu / array_lambda + + dfdict[dirbin] = pd.DataFrame( + { + "lambda_angstroms": array_lambda, + "f_lambda": array_flambda, + } + ) - if getpacketcount: - dfdict["packetcount"] = array_pktcount + if getpacketcount: + dfdict[dirbin]["packetcount"] = dfbinned["count"] - return pd.DataFrame(dfdict) + return dfdict @lru_cache(maxsize=16) def read_spec_res(modelpath: Path) -> dict[int, pd.DataFrame]: """Return dataframe of time-series spectra for every viewing direction""" - if Path(modelpath).is_file(): - specfilename = modelpath - else: - specfilename = at.firstexisting( - [ - "specpol_res.out", - "specpol_res.out.xz", - "specpol_res.out.gz", - "spec_res.out", - "spec_res.out.xz", - "spec_res.out.gz", - ], - path=modelpath, - ) + specfilename = ( + modelpath + if Path(modelpath).is_file() + else at.firstexisting(["spec_res.out", "specpol_res.out"], folder=modelpath, tryzipped=True) + ) print(f"Reading {specfilename} (in read_spec_res)") - specdata = pd.read_csv(specfilename, delim_whitespace=True, header=None, dtype=str) - - res_specdata: dict[int, pd.DataFrame] = at.gather_res_data(specdata) - - # index_to_split = specdata.index[specdata.iloc[:, 1] == specdata.iloc[0, 1]] - # # print(len(index_to_split)) - # res_specdata = [] - # for i, index_value in enumerate(index_to_split): - # if index_value != index_to_split[-1]: - # chunk = specdata.iloc[index_to_split[i]:index_to_split[i + 1], :] - # else: - # chunk = specdata.iloc[index_to_split[i]:, :] - # res_specdata.append(chunk) - # print(res_specdata[0]) - - columns = res_specdata[0].iloc[0] - # print(columns) - for i in res_specdata.keys(): - res_specdata[i] = res_specdata[i].rename(columns=columns) - res_specdata[i].drop(res_specdata[i].index[0], inplace=True) - # These lines remove the Q and U values from the dataframe (I think) - numberofIvalues = len(res_specdata[i].columns.drop_duplicates()) - res_specdata_numpy = res_specdata[i].iloc[:, :numberofIvalues].astype(float).to_numpy() - - res_specdata[i] = pd.DataFrame(data=res_specdata_numpy, columns=columns[:numberofIvalues]) - res_specdata[i].rename(columns={"0": "nu", "0.0": "nu"}, inplace=True) + res_specdata_in = pl.read_csv(at.zopen(specfilename, "rb"), separator=" ", has_header=False, infer_schema_length=0) + + # drop last column of nulls (caused by trailing space on each line) + if res_specdata_in[res_specdata_in.columns[-1]].is_null().all(): + res_specdata_in = res_specdata_in.drop(res_specdata_in.columns[-1]) + + res_specdata: dict[int, pl.DataFrame] = at.split_dataframe_dirbins(res_specdata_in, output_polarsdf=True) + prev_dfshape = None + for dirbin in res_specdata: + newcolnames = [str(x) for x in res_specdata[dirbin][0, :].to_numpy()[0]] + newcolnames[0] = "nu" + + newcolnames_unique = set(newcolnames) + oldcolnames = res_specdata[dirbin].columns + if len(newcolnames) > len(newcolnames_unique): + # for POL_ON, the time columns repeat for Q, U, and V stokes params. + # here, we keep the first set (I) and drop the rest of the columns + assert len(newcolnames) % len(newcolnames_unique) == 0 # must be an exact multiple + newcolnames = newcolnames[: len(newcolnames_unique)] + oldcolnames = oldcolnames[: len(newcolnames_unique)] + res_specdata[dirbin] = res_specdata[dirbin].select(oldcolnames) + + res_specdata[dirbin] = ( + res_specdata[dirbin][1:] # drop the first row that contains time headers + .with_columns(pl.all().cast(pl.Float64)) + .rename(dict(zip(oldcolnames, newcolnames))) + ) + + # the number of timesteps and nu bins should match for all direction bins + assert prev_dfshape is None or prev_dfshape == res_specdata[dirbin].shape + prev_dfshape = res_specdata[dirbin].shape return res_specdata -def average_phi_bins( - res_specdata: dict[int, pd.DataFrame], - dirbin: Optional[int], -) -> dict[int, pd.DataFrame]: - # Averages over phi (azimuthal) angle bins to make polar angle bin with less noise - dirbincount = at.get_viewingdirectionbincount() - phibincount = at.get_viewingdirection_phibincount() - assert dirbin is None or dirbin % phibincount == 0 - for start_bin in range(0, dirbincount, phibincount): - if dirbin is not None and start_bin != dirbin: - continue - res_specdata[start_bin] = res_specdata[start_bin].copy() # important to not affect the LRU cached copy - for bin_number in range(start_bin + 1, start_bin + phibincount): - res_specdata[start_bin] += res_specdata[bin_number] - del res_specdata[bin_number] - res_specdata[start_bin] /= phibincount - print( - f"Bin number {dirbin} is the average of {phibincount} bins {start_bin} to" f" {start_bin + phibincount - 1}" - ) +@lru_cache(maxsize=200) +def read_emission_absorption_file(emabsfilename: Union[str, Path]) -> pl.DataFrame: + """Read into a DataFrame one of: emission.out. emissionpol.out, emissiontrue.out, absorption.out.""" + try: + emissionfilesize = Path(emabsfilename).stat().st_size / 1024 / 1024 + print(f" Reading {emabsfilename} ({emissionfilesize:.2f} MiB)") - return res_specdata + except AttributeError: + print(f" Reading {emabsfilename}") + dfemabs = pl.read_csv( + at.zopen(emabsfilename, "rb").read(), separator=" ", has_header=False, infer_schema_length=0 + ).with_columns(pl.all().cast(pl.Float32, strict=False)) -def average_costheta_bins( - res_specdata: dict[int, pd.DataFrame], - dirbin: Optional[int], -) -> dict[int, pd.DataFrame]: - # Averages over cos theta (polar) angle bins to make azimuthal angle bins with less noise - dirbincount = at.get_viewingdirectionbincount() - nphibins = at.get_viewingdirection_phibincount() - ncosthetabins = at.get_viewingdirection_costhetabincount() - assert dirbin is None or dirbin < nphibins - for start_bin in range(0, nphibins): - if dirbin is not None and start_bin != dirbin: - continue - contribbins = range(start_bin + ncosthetabins, dirbincount, ncosthetabins) + # drop last column of nulls (caused by trailing space on each line) + if dfemabs[dfemabs.columns[-1]].is_null().all(): + dfemabs = dfemabs.drop(dfemabs.columns[-1]) - res_specdata[start_bin] = res_specdata[start_bin].copy() # important to not affect the LRU cached copy - for bin_number in contribbins: - res_specdata[start_bin] += res_specdata[bin_number] - del res_specdata[bin_number] - res_specdata[start_bin] /= nphibins + return dfemabs - print(f"bin number {start_bin} = the average of bins {[start_bin] + list(contribbins)}") + +@lru_cache(maxsize=4) +def get_spec_res( + modelpath: Path, + average_over_theta: bool = False, + average_over_phi: bool = False, +) -> dict[int, pd.DataFrame]: + res_specdata = read_spec_res(modelpath) + if average_over_theta: + res_specdata = at.average_direction_bins(res_specdata, overangle="theta") + if average_over_phi: + res_specdata = at.average_direction_bins(res_specdata, overangle="phi") return res_specdata -def get_res_spectrum( +def get_spectrum( modelpath: Path, timestepmin: int, - timestepmax: int = -1, - angle: Optional[int] = None, - res_specdata: Optional[dict[int, pd.DataFrame]] = None, + timestepmax: Optional[int] = None, + directionbins: Optional[Sequence[int]] = None, fnufilterfunc: Optional[Callable[[np.ndarray], np.ndarray]] = None, - args: Optional[argparse.Namespace] = None, -) -> pd.DataFrame: + average_over_theta: bool = False, + average_over_phi: bool = False, + stokesparam: Literal["I", "Q", "U"] = "I", +) -> dict[int, pd.DataFrame]: """Return a pandas DataFrame containing an ARTIS emergent spectrum.""" - if timestepmax < 0: + if timestepmax is None or timestepmax < 0: timestepmax = timestepmin - # print(f"Reading spectrum at timestep {timestepmin}") + if directionbins is None: + directionbins = [-1] + # keys are direction bins (or -1 for spherical average) + specdata: dict[int, pd.DataFrame] = {} - if angle is None: - assert args is not None - angle = args.plotviewingangle[0] - print("WARNING: no viewing direction specified. Using direction bin {angle}") + if -1 in directionbins: + # spherically averaged spectra + if stokesparam == "I": + try: + specfilename = at.firstexisting("spec.out", folder=modelpath, tryzipped=True) - if res_specdata is None: - res_specdata = read_spec_res(modelpath).copy() - if args and args.average_over_phi_angle: - res_specdata = average_phi_bins(res_specdata, angle) - if args and args.average_over_theta_angle: - res_specdata = average_costheta_bins(res_specdata, angle) + print(f"Reading {specfilename}") - nu = res_specdata[angle].loc[:, "nu"].values - arr_tmid = at.get_timestep_times_float(modelpath, loc="mid") - arr_tdelta = at.get_timestep_times_float(modelpath, loc="delta") + specdata[-1] = ( + pl.read_csv(at.zopen(specfilename, mode="rb"), separator=" ", infer_schema_length=0) + .with_columns(pl.all().cast(pl.Float64)) + .rename({"0": "nu"}) + ) - # for angle in args.plotviewingangle: - f_nu = stackspectra( - [ - (res_specdata[angle][res_specdata[angle].columns[timestep + 1]], arr_tdelta[timestep]) - for timestep in range(timestepmin, timestepmax + 1) - ] - ) + except FileNotFoundError: + specdata[-1] = get_specpol_data(angle=-1, modelpath=modelpath)[stokesparam] - # best to use the filter on this list because it - # has regular sampling - if fnufilterfunc: - print("Applying filter to ARTIS spectrum") - f_nu = fnufilterfunc(f_nu) + else: + specdata[-1] = get_specpol_data(angle=-1, modelpath=modelpath)[stokesparam] - dfspectrum = pd.DataFrame({"nu": nu, "f_nu": f_nu}) - dfspectrum.sort_values(by="nu", ascending=False, inplace=True) + if any(dirbin != -1 for dirbin in directionbins): + assert stokesparam == "I" + specdata.update( + get_spec_res(modelpath=modelpath, average_over_theta=average_over_theta, average_over_phi=average_over_phi) + ) - dfspectrum.eval("lambda_angstroms = @c / nu", local_dict={"c": 2.99792458e18}, inplace=True) - dfspectrum.eval("f_lambda = f_nu * nu / lambda_angstroms", inplace=True) - return dfspectrum + specdataout: dict[int, pd.DataFrame] = {} + for dirbin in directionbins: + arr_nu = specdata[dirbin]["nu"].to_numpy() + arr_tdelta = at.get_timestep_times_float(modelpath, loc="delta") + + arr_f_nu = stackspectra( + [ + (specdata[dirbin][specdata[dirbin].columns[timestep + 1]].to_numpy(), arr_tdelta[timestep]) + for timestep in range(timestepmin, timestepmax + 1) + ] + ) + + # best to use the filter on this list because it + # has regular sampling + if fnufilterfunc: + if dirbin == directionbins[0]: + print("Applying filter to ARTIS spectrum") + arr_f_nu = fnufilterfunc(arr_f_nu) + c_ang_per_s = 2.99792458e18 + arr_lambda = c_ang_per_s / arr_nu + arr_f_lambda = arr_f_nu * arr_nu / arr_lambda + dfspectrum = pd.DataFrame({"lambda_angstroms": arr_lambda, "f_lambda": arr_f_lambda}) + dfspectrum = dfspectrum.sort_values(by="lambda_angstroms", ascending=True) -def make_virtual_spectra_summed_file(modelpath: Path) -> None: + specdataout[dirbin] = dfspectrum + + return specdataout + + +def make_virtual_spectra_summed_file(modelpath: Path) -> Path: nprocs = at.get_nprocs(modelpath) print("nprocs", nprocs) - vspecpol_data_old: list[ - pd.DataFrame - ] = [] # virtual packet spectra for each observer (all directions and opacity choices) + vspecpol_data_old: list[pd.DataFrame] = ( + [] + ) # virtual packet spectra for each observer (all directions and opacity choices) vpktconfig = at.get_vpkt_config(modelpath) nvirtual_spectra = vpktconfig["nobsdirections"] * vpktconfig["nspectraperobs"] print( @@ -485,11 +430,12 @@ def make_virtual_spectra_summed_file(modelpath: Path) -> None: index_of_new_spectrum = vspecpolfile.index[vspecpolfile.iloc[:, 1] == vspecpolfile.iloc[0, 1]] vspecpol_data = [] # list of all predefined vspectra for i, index_spectrum_starts in enumerate(index_of_new_spectrum[:nvirtual_spectra]): - # todo: this is different to at.gather_res_data() -- could be made to be same format to not repeat code - if index_spectrum_starts != index_of_new_spectrum[-1]: - chunk = vspecpolfile.iloc[index_spectrum_starts : index_of_new_spectrum[i + 1], :] - else: - chunk = vspecpolfile.iloc[index_spectrum_starts:, :] + # todo: this is different to at.split_dataframe_dirbins() -- could be made to be same format to not repeat code + chunk = ( + vspecpolfile.iloc[index_spectrum_starts : index_of_new_spectrum[i + 1], :] + if index_spectrum_starts != index_of_new_spectrum[-1] + else vspecpolfile.iloc[index_spectrum_starts:, :] + ) vspecpol_data.append(chunk) if len(vspecpol_data_old) > 0: @@ -508,6 +454,8 @@ def make_virtual_spectra_summed_file(modelpath: Path) -> None: print(f"Saved {outfile}") vspecpol.to_csv(outfile, sep=" ", index=False, header=False) + return outfile + def make_averaged_vspecfiles(args: argparse.Namespace) -> None: filenames = [] @@ -539,43 +487,75 @@ def alphanum_key(key: str) -> list[Union[int, str]]: ) +@lru_cache(maxsize=4) def get_specpol_data( - angle: Optional[int] = None, modelpath: Optional[Path] = None, specdata: Optional[pd.DataFrame] = None + angle: int = -1, modelpath: Optional[Path] = None, specdata: Optional[pd.DataFrame] = None ) -> dict[str, pd.DataFrame]: if specdata is None: assert modelpath is not None - if angle is None: - specfilename = at.firstexisting("specpol.out", path=modelpath, tryzipped=True) - else: - # alternatively use f'vspecpol_averaged-{angle}.out' ? - vspecpath = modelpath - if os.path.isdir(modelpath / "vspecpol"): - vspecpath = modelpath / "vspecpol" - specfilename = at.firstexisting(f"vspecpol_total-{angle}.out", path=vspecpath, tryzipped=True) - if not specfilename.exists(): - print(f"{specfilename} does not exist. Generating all-rank summed vspec files..") - make_virtual_spectra_summed_file(modelpath=modelpath) + specfilename = ( + at.firstexisting("specpol.out", folder=modelpath, tryzipped=True) + if angle == -1 + else at.firstexisting(f"specpol_res_{angle}.out", folder=modelpath, tryzipped=True) + ) + + print(f"Reading {specfilename}") + specdata = pd.read_csv(specfilename, delim_whitespace=True) + + stokes_params = split_dataframe_stokesparams(specdata) + + return stokes_params + + +@lru_cache(maxsize=4) +def get_vspecpol_data( + vspecangle: Optional[int] = None, modelpath: Optional[Path] = None, specdata: Optional[pd.DataFrame] = None +) -> dict[str, pd.DataFrame]: + if specdata is None: + assert modelpath is not None + # alternatively use f'vspecpol_averaged-{angle}.out' ? + vspecpath = modelpath + if os.path.isdir(modelpath / "vspecpol"): + vspecpath = modelpath / "vspecpol" + + try: + specfilename = at.firstexisting(f"vspecpol_total-{vspecangle}.out", folder=vspecpath, tryzipped=True) + except FileNotFoundError: + print(f"vspecpol_total-{vspecangle}.out does not exist. Generating all-rank summed vspec files..") + specfilename = make_virtual_spectra_summed_file(modelpath=modelpath) print(f"Reading {specfilename}") specdata = pd.read_csv(specfilename, delim_whitespace=True) - specdata = specdata.rename(columns={specdata.keys()[0]: "nu"}) + stokes_params = split_dataframe_stokesparams(specdata) + + return stokes_params + + +def split_dataframe_stokesparams(specdata: pd.DataFrame) -> dict[str, pd.DataFrame]: + """DataFrames read from specpol*.out and vspecpol*.out are repeated over I, Q, U + parameters. Split these into a dictionary of DataFrames. + """ + specdata = specdata.rename({"0": "nu", "0.0": "nu"}, axis="columns") cols_to_split = [] stokes_params = {} for i, key in enumerate(specdata.keys()): if specdata.keys()[1] in key: cols_to_split.append(i) - stokes_params["I"] = pd.concat([specdata["nu"], specdata.iloc[:, cols_to_split[0] : cols_to_split[1]]], axis=1) - stokes_params["Q"] = pd.concat([specdata["nu"], specdata.iloc[:, cols_to_split[1] : cols_to_split[2]]], axis=1) - stokes_params["U"] = pd.concat([specdata["nu"], specdata.iloc[:, cols_to_split[2] :]], axis=1) + stokes_params["I"] = pd.concat( + [specdata["nu"], specdata.iloc[:, cols_to_split[0] : cols_to_split[1]]], axis="columns" + ) + stokes_params["Q"] = pd.concat( + [specdata["nu"], specdata.iloc[:, cols_to_split[1] : cols_to_split[2]]], axis="columns" + ) + stokes_params["U"] = pd.concat([specdata["nu"], specdata.iloc[:, cols_to_split[2] :]], axis="columns") for param in ["Q", "U"]: stokes_params[param].columns = stokes_params["I"].keys() stokes_params[param + "/I"] = pd.concat( - [specdata["nu"], stokes_params[param].iloc[:, 1:] / stokes_params["I"].iloc[:, 1:]], axis=1 + [specdata["nu"], stokes_params[param].iloc[:, 1:] / stokes_params["I"].iloc[:, 1:]], axis="columns" ) - return stokes_params @@ -586,14 +566,14 @@ def get_vspecpol_spectrum( args: argparse.Namespace, fnufilterfunc: Optional[Callable[[np.ndarray], np.ndarray]] = None, ) -> pd.DataFrame: - stokes_params = get_specpol_data(angle, modelpath=Path(modelpath)) + stokes_params = get_vspecpol_data(vspecangle=angle, modelpath=Path(modelpath)) if "stokesparam" not in args: args.stokesparam = "I" vspecdata = stokes_params[args.stokesparam] - nu = vspecdata.loc[:, "nu"].values + nu = vspecdata.loc[:, "nu"].to_numpy() - arr_tmid = [float(i) for i in vspecdata.columns.values[1:] if i[-2] != "."] + arr_tmid = [float(i) for i in vspecdata.columns.to_numpy()[1:] if i[-2] != "."] arr_tdelta = [l1 - l2 for l1, l2 in zip(arr_tmid[1:], arr_tmid[:-1])] + [arr_tmid[-1] - arr_tmid[-2]] def match_closest_time(reftime: float) -> str: @@ -622,10 +602,10 @@ def match_closest_time(reftime: float) -> str: f_nu = fnufilterfunc(f_nu) dfspectrum = pd.DataFrame({"nu": nu, "f_nu": f_nu}) - dfspectrum.sort_values(by="nu", ascending=False, inplace=True) + dfspectrum = dfspectrum.sort_values(by="nu", ascending=False) - dfspectrum.eval("lambda_angstroms = @c / nu", local_dict={"c": 2.99792458e18}, inplace=True) - dfspectrum.eval("f_lambda = f_nu * nu / lambda_angstroms", inplace=True) + dfspectrum = dfspectrum.eval("lambda_angstroms = @c / nu", local_dict={"c": 2.99792458e18}) + dfspectrum = dfspectrum.eval("f_lambda = f_nu * nu / lambda_angstroms") return dfspectrum @@ -645,7 +625,7 @@ def get_flux_contributions( ) -> tuple[list[fluxcontributiontuple], np.ndarray]: arr_tmid = at.get_timestep_times_float(modelpath, loc="mid") arr_tdelta = at.get_timestep_times_float(modelpath, loc="delta") - arraynu = at.misc.get_nu_grid(modelpath) + arraynu = at.get_nu_grid(modelpath) arraylambda = 2.99792458e18 / arraynu if not Path(modelpath, "compositiondata.txt").is_file(): print("WARNING: compositiondata.txt not found. Using output*.txt instead") @@ -658,11 +638,11 @@ def get_flux_contributions( dbinlist = [-1] elif averageoverphi: assert not averageovertheta - assert directionbin is not None and directionbin % at.get_viewingdirection_phibincount() == 0 + assert directionbin % at.get_viewingdirection_phibincount() == 0 dbinlist = list(range(directionbin, directionbin + at.get_viewingdirection_phibincount())) elif averageovertheta: assert not averageoverphi - assert directionbin is not None and directionbin < at.get_viewingdirection_phibincount() + assert directionbin < at.get_viewingdirection_phibincount() dbinlist = list(range(directionbin, at.get_viewingdirectionbincount(), at.get_viewingdirection_phibincount())) else: dbinlist = [directionbin] @@ -670,37 +650,21 @@ def get_flux_contributions( emissiondata: dict[int, pd.DataFrame] = {} absorptiondata: dict[int, pd.DataFrame] = {} maxion: Optional[int] = None - for i, dbin in enumerate(dbinlist): + for dbin in dbinlist: if getemission: - if use_lastemissiontype: - emissionfilenames = ["emission.out", "emissionpol.out"] - else: - emissionfilenames = ["emissiontrue.out"] + emissionfilenames = ["emission.out", "emissionpol.out"] if use_lastemissiontype else ["emissiontrue.out"] if dbin != -1: emissionfilenames = [x.replace(".out", f"_res_{dbin:02d}.out") for x in emissionfilenames] - emissionfilename = at.firstexisting(emissionfilenames, path=modelpath, tryzipped=True) + emissionfilename = at.firstexisting(emissionfilenames, folder=modelpath, tryzipped=True) if "pol" in str(emissionfilename): print("This artis run contains polarisation data") # File contains I, Q and U and so times are repeated 3 times arr_tmid = np.array(arr_tmid.tolist() * 3) - try: - emissionfilesize = Path(emissionfilename).stat().st_size / 1024 / 1024 - print(f" Reading {emissionfilename} ({emissionfilesize:.2f} MiB)") - - except AttributeError: - print(f" Reading {emissionfilename}") - - emissiondata[dbin] = pd.read_table( - emissionfilename, sep=" ", engine=at.get_config()["pandas_engine"], header=None - ) - - # check if last column is an artefact of whitespace at end of line (None or NaNs for pyarrow/c engine) - if emissiondata[dbin].iloc[0, -1] is None or np.isnan(emissiondata[dbin].iloc[0, -1]): - emissiondata[dbin].drop(emissiondata[dbin].columns[-1], axis=1, inplace=True) + emissiondata[dbin] = read_emission_absorption_file(emissionfilename) maxion_float = (emissiondata[dbin].shape[1] - 1) / 2 / nelements # also known as MIONS in ARTIS sn3d.h assert maxion_float.is_integer() @@ -721,22 +685,9 @@ def get_flux_contributions( if directionbin is not None: absorptionfilenames = [x.replace(".out", f"_res_{dbin:02d}.out") for x in absorptionfilenames] - absorptionfilename = at.firstexisting(absorptionfilenames, path=modelpath, tryzipped=True) - - try: - absorptionfilesize = Path(absorptionfilename).stat().st_size / 1024 / 1024 - print(f" Reading {absorptionfilename} ({absorptionfilesize:.2f} MiB)") - except AttributeError: - print(f" Reading {absorptionfilename}") - - absorptiondata[dbin] = pd.read_table( - absorptionfilename, sep=" ", engine=at.get_config()["pandas_engine"], header=None - ) - - # check if last column is an artefact of whitespace at end of line (None or NaNs for pyarrow/c engine) - if absorptiondata[dbin].iloc[0, -1] is None or np.isnan(absorptiondata[dbin].iloc[0, -1]): - absorptiondata[dbin].drop(absorptiondata[dbin].columns[-1], axis=1, inplace=True) + absorptionfilename = at.firstexisting(absorptionfilenames, folder=modelpath, tryzipped=True) + absorptiondata[dbin] = read_emission_absorption_file(absorptionfilename) absorption_maxion_float = absorptiondata[dbin].shape[1] / nelements assert absorption_maxion_float.is_integer() absorption_maxion = int(absorption_maxion_float) @@ -776,7 +727,7 @@ def get_flux_contributions( array_fnu_emission = stackspectra( [ ( - emissiondata[dbin].iloc[timestep :: len(arr_tmid), selectedcolumn].values, + emissiondata[dbin][timestep :: len(arr_tmid), selectedcolumn].to_numpy(), arr_tdelta[timestep] / len(dbinlist), ) for timestep in range(timestepmin, timestepmax + 1) @@ -790,7 +741,7 @@ def get_flux_contributions( array_fnu_absorption = stackspectra( [ ( - absorptiondata[dbin].iloc[timestep :: len(arr_tmid), selectedcolumn].values, + absorptiondata[dbin][timestep :: len(arr_tmid), selectedcolumn].to_numpy(), arr_tdelta[timestep] / len(dbinlist), ) for timestep in range(timestepmin, timestepmax + 1) @@ -864,14 +815,12 @@ def get_emprocesslabel( if emtype >= 0: line = linelist[emtype] if groupby == "line": - # if line.atomic_number != 26 or line.ionstage != 2: - # return 'non-Fe II ions' return ( f"{at.get_ionstring(line.atomic_number, line.ionstage)} " f"λ{line.lambda_angstroms:.0f} " f"({line.upperlevelindex}-{line.lowerlevelindex})" ) - elif groupby == "terms": + if groupby == "terms": upper_config = ( adata.query("Z == @line.atomic_number and ion_stage == @line.ionstage", inplace=False) .iloc[0] @@ -887,7 +836,7 @@ def get_emprocesslabel( ) lower_term_noj = lower_config.split("_")[-1].split("[")[0] return f"{at.get_ionstring(line.atomic_number, line.ionstage)} {upper_term_noj}->{lower_term_noj}" - elif groupby == "upperterm": + if groupby == "upperterm": upper_config = ( adata.query("Z == @line.atomic_number and ion_stage == @line.ionstage", inplace=False) .iloc[0] @@ -897,7 +846,7 @@ def get_emprocesslabel( upper_term_noj = upper_config.split("_")[-1].split("[")[0] return f"{at.get_ionstring(line.atomic_number, line.ionstage)} {upper_term_noj}" return f"{at.get_ionstring(line.atomic_number, line.ionstage)} bound-bound" - elif emtype == -9999999: + if emtype == -9999999: return "free-free" bfindex = -emtype - 1 @@ -937,8 +886,6 @@ def get_absprocesslabel(linelist: dict[int, at.linetuple], abstype: int) -> str: vmax = modeldata.iloc[-1].velocity_outer * u.km / u.s betafactor = math.sqrt(1 - (vmax / const.c).decompose().value ** 2) - import artistools.packets - packetsfiles = at.packets.get_packetsfilepaths(modelpath, maxpacketfiles) linelist = at.get_linelist_dict(modelpath=modelpath) @@ -954,12 +901,11 @@ def get_absprocesslabel(linelist: dict[int, at.linetuple], abstype: int) -> str: nu_min = 2.99792458e18 / lambda_max nu_max = 2.99792458e18 / lambda_min - if useinternalpackets: - emtypecolumn = "emissiontype" - else: - emtypecolumn = "emissiontype" if use_lastemissiontype else "trueemissiontype" + emtypecolumn = ( + "emissiontype" if useinternalpackets else "emissiontype" if use_lastemissiontype else "trueemissiontype" + ) - for index, packetsfile in enumerate(packetsfiles): + for _index, packetsfile in enumerate(packetsfiles): if useinternalpackets: # if we're using packets*.out files, these packets are from the last timestep t_seconds = at.get_timestep_times_float(modelpath, loc="start")[-1] * 86400.0 @@ -974,42 +920,38 @@ def get_absprocesslabel(linelist: dict[int, at.linetuple], abstype: int) -> str: r_inner = t_seconds * v_inner r_outer = t_seconds * v_outer - dfpackets = at.packets.readfile(packetsfile, type="TYPE_RPKT") + dfpackets = at.packets.readfile(packetsfile, packet_type="TYPE_RPKT") print("Using non-escaped internal r-packets") - dfpackets.query( - f'type_id == {at.packets.type_ids["TYPE_RPKT"]} and @nu_min <= nu_rf < @nu_max', inplace=True - ) + dfpackets = dfpackets.query(f'type_id == {at.packets.type_ids["TYPE_RPKT"]} and @nu_min <= nu_rf < @nu_max') if modelgridindex is not None: assoc_cells, mgi_of_propcells = at.get_grid_mapping(modelpath=modelpath) # dfpackets.eval(f'velocity = sqrt(posx ** 2 + posy ** 2 + posz ** 2) / @t_seconds', inplace=True) # dfpackets.query(f'@v_inner <= velocity <= @v_outer', # inplace=True) - dfpackets.query("where in @assoc_cells[@modelgridindex]", inplace=True) + dfpackets = dfpackets.query("where in @assoc_cells[@modelgridindex]") print(f" {len(dfpackets)} internal r-packets matching frequency range") else: - dfpackets = at.packets.readfile(packetsfile, type="TYPE_ESCAPE", escape_type="TYPE_RPKT") - dfpackets.query( - "@nu_min <= nu_rf < @nu_max and " + dfpackets = at.packets.readfile(packetsfile, packet_type="TYPE_ESCAPE", escape_type="TYPE_RPKT") + dfpackets = dfpackets.query( + "@nu_min <= nu_rf < @nu_max and trueemissiontype >= 0 and " + ( "@timelow < (escape_time - (posx * dirx + posy * diry + posz * dirz) / @c_cgs) < @timehigh" if not use_escapetime else "@timelow < escape_time * @betafactor < @timehigh" ), - inplace=True, ) print(f" {len(dfpackets)} escaped r-packets matching frequency and arrival time ranges") if emissionvelocitycut: dfpackets = at.packets.add_derived_columns(dfpackets, modelpath, ["emission_velocity"]) - dfpackets.query("(emission_velocity / 1e5) > @emissionvelocitycut", inplace=True) + dfpackets = dfpackets.query("(emission_velocity / 1e5) > @emissionvelocitycut") if np.isscalar(delta_lambda): - dfpackets.eval("xindex = floor((2.99792458e18 / nu_rf - @lambda_min) / @delta_lambda)", inplace=True) + dfpackets = dfpackets.eval("xindex = floor((2.99792458e18 / nu_rf - @lambda_min) / @delta_lambda)") if getabsorption: - dfpackets.eval( + dfpackets = dfpackets.eval( "xindexabsorbed = floor((2.99792458e18 / absorption_freq - @lambda_min) / @delta_lambda)", - inplace=True, ) else: dfpackets["xindex"] = ( @@ -1022,7 +964,6 @@ def get_absprocesslabel(linelist: dict[int, at.linetuple], abstype: int) -> str: bflist = at.get_bflist(modelpath) for _, packet in dfpackets.iterrows(): - lambda_rf = 2.99792458e18 / packet.nu_rf xindex = int(packet.xindex) assert xindex >= 0 @@ -1137,13 +1078,13 @@ def sortkey(x: fluxcontributiontuple) -> tuple[int, float]: remainder_fluxcontrib = 0 if greyscale: - hatches = artistools.spectra.plotspectra.hatches + hatches = at.spectra.plotspectra.hatches seriescount = len(fixedionlist) if fixedionlist else maxseriescount colorcount = math.ceil(seriescount / 1.0 / len(hatches)) greylist = [str(x) for x in np.linspace(0.4, 0.9, colorcount, endpoint=True)] color_list = [] for c in range(colorcount): - for h in hatches: + for _h in hatches: color_list.append(greylist[c]) # color_list = list(plt.get_cmap('tab20')(np.linspace(0, 1.0, 20))) mpl.rcParams["hatch.linewidth"] = 0.1 @@ -1208,7 +1149,7 @@ def sortkey(x: fluxcontributiontuple) -> tuple[int, float]: def print_integrated_flux( arr_f_lambda: np.ndarray, arr_lambda_angstroms: np.ndarray, distance_megaparsec: float = 1.0 -) -> None: +) -> float: integrated_flux = abs(np.trapz(arr_f_lambda, x=arr_lambda_angstroms)) * u.erg / u.s / (u.cm**2) print( f" integrated flux ({arr_lambda_angstroms.min():.1f} to " @@ -1216,14 +1157,15 @@ def print_integrated_flux( ) # luminosity = integrated_flux * 4 * math.pi * (distance_megaparsec * u.megaparsec ** 2) # print(f'(L={luminosity.to("Lsun"):.3e})') + return integrated_flux def get_line_flux( lambda_low: float, lambda_high: float, arr_f_lambda: np.ndarray, arr_lambda_angstroms: np.ndarray ) -> float: - index_low, index_high = [ + index_low, index_high = ( int(np.searchsorted(arr_lambda_angstroms, wl, side="left")) for wl in (lambda_low, lambda_high) - ] + ) flux_integral = abs(np.trapz(arr_f_lambda[index_low:index_high], x=arr_lambda_angstroms[index_low:index_high])) return flux_integral @@ -1258,7 +1200,7 @@ def get_reference_spectrum(filename: Union[Path, str]) -> tuple[pd.DataFrame, di if filepathgz.is_file(): filepath = filepathgz - metadata = at.misc.get_file_metadata(filepath) + metadata = at.get_file_metadata(filepath) flambdaindex = metadata.get("f_lambda_columnindex", 1) @@ -1299,11 +1241,12 @@ def get_reference_spectrum(filename: Union[Path, str]) -> tuple[pd.DataFrame, di if "a_v" in metadata or "e_bminusv" in metadata: print("Correcting for reddening") - from extinction import apply, ccm89 + from extinction import apply + from extinction import ccm89 specdata["f_lambda"] = apply( - ccm89(specdata["lambda_angstroms"].values, a_v=-metadata["a_v"], r_v=metadata["r_v"], unit="aa"), - specdata["f_lambda"].values, + ccm89(specdata["lambda_angstroms"].to_numpy(), a_v=-metadata["a_v"], r_v=metadata["r_v"], unit="aa"), + specdata["f_lambda"].to_numpy(), ) if "z" in metadata: @@ -1326,11 +1269,11 @@ def write_flambda_spectra(modelpath: Path, args: argparse.Namespace) -> None: if Path(modelpath, "specpol.out").is_file(): specfilename = modelpath / "specpol.out" specdata = pd.read_csv(specfilename, delim_whitespace=True) - timearray = [i for i in specdata.columns.values[1:] if i[-2] != "."] + timearray = [i for i in specdata.columns.to_numpy()[1:] if i[-2] != "."] else: - specfilename = at.firstexisting("spec.out", path=modelpath, tryzipped=True) + specfilename = at.firstexisting("spec.out", folder=modelpath, tryzipped=True) specdata = pd.read_csv(specfilename, delim_whitespace=True) - timearray = specdata.columns.values[1:] + timearray = specdata.columns.to_numpy()[1:] number_of_timesteps = len(timearray) @@ -1345,7 +1288,7 @@ def write_flambda_spectra(modelpath: Path, args: argparse.Namespace) -> None: arr_tmid = at.get_timestep_times_float(modelpath, loc="mid") for timestep in range(timestepmin, timestepmax + 1): - dfspectrum = get_spectrum(modelpath, timestep, timestep) + dfspectrum = get_spectrum(modelpath=modelpath, timestepmin=timestep, timestepmax=timestep)[-1] tmid = arr_tmid[timestep] outfilepath = outdirectory / f"spectrum_ts{timestep:02.0f}_{tmid:.2f}d.txt" diff --git a/artistools/spectra/test_spectra.py b/artistools/spectra/test_spectra.py index 84bba7a09..862f1f33a 100755 --- a/artistools/spectra/test_spectra.py +++ b/artistools/spectra/test_spectra.py @@ -1,23 +1,19 @@ #!/usr/bin/env python3 import math import os.path -from pathlib import Path import numpy as np import pandas as pd from astropy import constants as const import artistools as at -import artistools.spectra -import artistools.transitions modelpath = at.get_config()["path_testartismodel"] outputpath = at.get_config()["path_testoutput"] -at.set_config("enable_diskcache", False) -def test_spectraplot(): - at.spectra.main( +def test_spectraplot() -> None: + at.spectra.plot( argsraw=[], specpath=[modelpath, "sn2011fe_PTF11kly_20120822_norm.txt"], outputfile=outputpath, @@ -26,8 +22,8 @@ def test_spectraplot(): ) -def test_spectra_frompackets(): - at.spectra.main( +def test_spectra_frompackets() -> None: + at.spectra.plot( argsraw=[], specpath=modelpath, outputfile=os.path.join(outputpath, "spectrum_from_packets.pdf"), @@ -37,12 +33,12 @@ def test_spectra_frompackets(): ) -def test_spectra_outputtext(): - at.spectra.main(argsraw=[], specpath=modelpath, output_spectra=True) +def test_spectra_outputtext() -> None: + at.spectra.plot(argsraw=[], specpath=modelpath, output_spectra=True) -def test_spectraemissionplot(): - at.spectra.main( +def test_spectraemissionplot() -> None: + at.spectra.plot( argsraw=[], specpath=modelpath, outputfile=outputpath, @@ -53,8 +49,8 @@ def test_spectraemissionplot(): ) -def test_spectraemissionplot_nostack(): - at.spectra.main( +def test_spectraemissionplot_nostack() -> None: + at.spectra.plot( argsraw=[], specpath=modelpath, outputfile=outputpath, @@ -66,43 +62,43 @@ def test_spectraemissionplot_nostack(): ) -def test_spectra_get_spectrum(): - def check_spectrum(dfspectrumpkts): +def test_spectra_get_spectrum() -> None: + def check_spectrum(dfspectrumpkts) -> None: assert math.isclose(max(dfspectrumpkts["f_lambda"]), 2.548532804918824e-13, abs_tol=1e-5) assert min(dfspectrumpkts["f_lambda"]) < 1e-9 assert math.isclose(np.mean(dfspectrumpkts["f_lambda"]), 1.0314682640070206e-14, abs_tol=1e-5) - dfspectrum = at.spectra.get_spectrum(modelpath, 55, 65, fnufilterfunc=None) + dfspectrum = at.spectra.get_spectrum(modelpath, 55, 65, fnufilterfunc=None)[-1] assert len(dfspectrum["lambda_angstroms"]) == 1000 assert len(dfspectrum["f_lambda"]) == 1000 - assert abs(dfspectrum["lambda_angstroms"].values[-1] - 29920.601421214415) < 1e-5 - assert abs(dfspectrum["lambda_angstroms"].values[0] - 600.75759482509852) < 1e-5 + assert abs(dfspectrum["lambda_angstroms"].to_numpy()[-1] - 29920.601421214415) < 1e-5 + assert abs(dfspectrum["lambda_angstroms"].to_numpy()[0] - 600.75759482509852) < 1e-5 check_spectrum(dfspectrum) - lambda_min = dfspectrum["lambda_angstroms"].values[0] - lambda_max = dfspectrum["lambda_angstroms"].values[-1] + lambda_min = dfspectrum["lambda_angstroms"].to_numpy()[0] + lambda_max = dfspectrum["lambda_angstroms"].to_numpy()[-1] timelowdays = at.get_timestep_times_float(modelpath)[55] timehighdays = at.get_timestep_times_float(modelpath)[65] - dfspectrumpkts = at.spectra.get_spectrum_from_packets( + dfspectrumpkts = at.spectra.get_from_packets( modelpath, timelowdays=timelowdays, timehighdays=timehighdays, lambda_min=lambda_min, lambda_max=lambda_max - ) + )[-1] check_spectrum(dfspectrumpkts) -def test_spectra_get_flux_contributions(): +def test_spectra_get_flux_contributions() -> None: timestepmin = 40 timestepmax = 80 dfspectrum = at.spectra.get_spectrum( - modelpath, timestepmin=timestepmin, timestepmax=timestepmax, fnufilterfunc=None - ) + modelpath=modelpath, timestepmin=timestepmin, timestepmax=timestepmax, fnufilterfunc=None + )[-1] integrated_flux_specout = np.trapz(dfspectrum["f_lambda"], x=dfspectrum["lambda_angstroms"]) specdata = pd.read_csv(modelpath / "spec.out", delim_whitespace=True) - arraynu = specdata.loc[:, "0"].values + arraynu = specdata.loc[:, "0"].to_numpy() arraylambda_angstroms = const.c.to("angstrom/s").value / arraynu contribution_list, array_flambda_emission_total = at.spectra.get_flux_contributions( @@ -120,13 +116,13 @@ def test_spectra_get_flux_contributions(): assert math.isclose(integrated_flux_specout, integrated_flux_emission, rel_tol=4e-3) # check each bin is not out by a large fraction - diff = [abs(x - y) for x, y in zip(array_flambda_emission_total, dfspectrum["f_lambda"].values)] + diff = [abs(x - y) for x, y in zip(array_flambda_emission_total, dfspectrum["f_lambda"].to_numpy())] print(f"Max f_lambda difference {max(diff) / integrated_flux_specout}") assert max(diff) / integrated_flux_specout < 2e-3 -def test_spectra_timeseries_subplots(): +def test_spectra_timeseries_subplots() -> None: timedayslist = [295, 300] - at.spectra.main( + at.spectra.plot( argsraw=[], specpath=modelpath, outputfile=outputpath, timedayslist=timedayslist, multispecplot=True ) diff --git a/artistools/spectra/test_vspectra.py b/artistools/spectra/test_vspectra.py new file mode 100755 index 000000000..3bf27cf3c --- /dev/null +++ b/artistools/spectra/test_vspectra.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +from unittest import mock + +import matplotlib.axes +import numpy as np + +import artistools as at + +modelpath = at.get_config()["path_testartismodel"].parent / "vspecpolmodel" +outputpath = at.get_config()["path_testoutput"] + + +@mock.patch.object(matplotlib.axes.Axes, "plot", side_effect=matplotlib.axes.Axes.plot, autospec=True) +def test_vspectraplot(mockplot): + at.spectra.plot( + argsraw=[], + specpath=[modelpath, "sn2011fe_PTF11kly_20120822_norm.txt"], + outputfile=outputpath, + plotvspecpol=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + timemin=10, + timemax=12, + ) + + arr_time_d = np.array(mockplot.call_args_list[0][0][1]) + assert all(np.array_equal(arr_time_d, np.array(mockplot.call_args_list[vspecdir][0][1])) for vspecdir in range(10)) + + arr_allvspec = np.vstack([np.array(mockplot.call_args_list[vspecdir][0][2]) for vspecdir in range(10)]) + assert np.allclose( + arr_allvspec.std(axis=1), + np.array( + [ + 2.01529689e-12, + 2.05807110e-12, + 2.01551623e-12, + 2.18216916e-12, + 2.85477069e-12, + 3.34384407e-12, + 2.94892344e-12, + 2.29084411e-12, + 2.05916843e-12, + 2.00515984e-12, + ] + ), + ) diff --git a/artistools/stats.py b/artistools/stats.py index a6fb34ef0..458018276 100755 --- a/artistools/stats.py +++ b/artistools/stats.py @@ -26,7 +26,7 @@ def main(): runfolder = logfilename.split("/output_0-0.txt")[0] timesteptimes = [] - with open(runfolder + "/light_curve.out", "r") as lcfile: + with open(runfolder + "/light_curve.out") as lcfile: for line in lcfile: timesteptimes.append(line.split()[0]) diff --git a/artistools/test_artistools.py b/artistools/test_artistools.py index fb488d59c..9b8e2a8dc 100755 --- a/artistools/test_artistools.py +++ b/artistools/test_artistools.py @@ -1,25 +1,20 @@ #!/usr/bin/env python3 import hashlib import math -import os.path -from pathlib import Path import numpy as np -import pandas as pd -import pytest import artistools as at modelpath = at.get_config()["path_testartismodel"] outputpath = at.get_config()["path_testoutput"] -at.set_config("enable_diskcache", False) def test_commands(): import importlib # ensure that the commands are pointing to valid submodule.function() targets - for command, (submodulename, funcname) in sorted(at.commands.get_commandlist().items()): + for _command, (submodulename, funcname) in sorted(at.commands.get_commandlist().items()): submodule = importlib.import_module(submodulename, package="artistools") assert hasattr(submodule, funcname) @@ -33,7 +28,7 @@ def test_timestep_times(): assert math.isclose(float(timemidarray[-1]), 349.412, abs_tol=1e-3) assert all( - [tstart < tmid < (tstart + tdelta) for tstart, tdelta, tmid in zip(timestartarray, timedeltarray, timemidarray)] + tstart < tmid < (tstart + tdelta) for tstart, tdelta, tmid in zip(timestartarray, timedeltarray, timemidarray) ) @@ -42,11 +37,11 @@ def test_deposition(): def test_estimator_snapshot(): - at.estimators.main(argsraw=[], modelpath=modelpath, outputfile=outputpath, timedays=300) + at.estimators.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, timedays=300) def test_estimator_timeevolution(): - at.estimators.main(argsraw=[], modelpath=modelpath, outputfile=outputpath, modelgridindex=0, x="time") + at.estimators.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, modelgridindex=0, x="time") def test_get_inputparams(): @@ -60,12 +55,6 @@ def test_get_levels(): def test_get_modeldata_tuple(): - # expect a 3D model but read 1D - with pytest.raises(Exception): - dfmodel, t_model_init_days, vmax_cmps = at.inputmodel.get_modeldata_tuple( - modelpath, get_elemabundances=True, dimensions=3 - ) - dfmodel, t_model_init_days, vmax_cmps = at.inputmodel.get_modeldata_tuple(modelpath, get_elemabundances=True) assert np.isclose(t_model_init_days, 0.00115740740741, rtol=0.0001) assert np.isclose(vmax_cmps, 800000000.0, rtol=0.0001) @@ -75,47 +64,18 @@ def test_get_modeldata_tuple(): # '40a02dfa933f6b28671d42f3cf69a182955a5a89dc93bbcd22c894192375fe9b') -def test_lightcurve(): - at.lightcurve.main(argsraw=[], modelpath=modelpath, outputfile=outputpath) - - -def test_lightcurve_frompackets(): - at.lightcurve.main( - argsraw=[], - modelpath=modelpath, - frompackets=True, - outputfile=os.path.join(outputpath, "lightcurve_from_packets.pdf"), - ) - - -def test_band_lightcurve_plot(): - at.lightcurve.main(argsraw=[], modelpath=modelpath, filter=["B"], outputfile=outputpath) - - -def test_band_lightcurve_subplots(): - at.lightcurve.main(argsraw=[], modelpath=modelpath, filter=["bol", "B"], outputfile=outputpath) - - -def test_colour_evolution_plot(): - at.lightcurve.main(argsraw=[], modelpath=modelpath, colour_evolution=["B-V"], outputfile=outputpath) - - -def test_colour_evolution_subplots(): - at.lightcurve.main(argsraw=[], modelpath=modelpath, colour_evolution=["U-B", "B-V"], outputfile=outputpath) - - def test_macroatom(): at.macroatom.main(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=10) def test_nltepops(): - # at.nltepops.main(modelpath=modelpath, outputfile=outputpath, timedays=300), + # at.nltepops.plot(modelpath=modelpath, outputfile=outputpath, timedays=300), # **benchargs) - at.nltepops.main(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=40) + at.nltepops.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=40) def test_nonthermal(): - at.nonthermal.main(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=70) + at.nonthermal.plot(argsraw=[], modelpath=modelpath, outputfile=outputpath, timestep=70) def test_radfield(): diff --git a/artistools/transitions.py b/artistools/transitions.py old mode 100644 new mode 100755 index 1a93f38db..255c9dd7d --- a/artistools/transitions.py +++ b/artistools/transitions.py @@ -2,6 +2,7 @@ import argparse import math import multiprocessing +import sys from collections import namedtuple from pathlib import Path @@ -11,14 +12,6 @@ from astropy import constants as const import artistools as at -import artistools.estimators -import artistools.nltepops -import artistools.spectra - -# import glob -# import re -# import numexpr as ne -# from astropy import units as u defaultoutputfile = "plottransitions_cell{cell:03d}_ts{timestep:02d}_{time_days:.0f}d.pdf" @@ -32,7 +25,7 @@ def get_kurucz_transitions(): ) translist = [] ionlist = [] - with open("gfall.dat", "r") as fnist: + with open("gfall.dat") as fnist: for line in fnist: row = line.split() if len(row) >= 24: @@ -61,7 +54,7 @@ def get_kurucz_transitions(): def get_nist_transitions(filename): transitiontuple = namedtuple("transition", "lambda_angstroms A lower_energy_ev upper_energy_ev lower_g upper_g") translist = [] - with open(filename, "r") as fnist: + with open(filename) as fnist: for line in fnist: row = line.split("|") if len(row) == 17 and "-" in row[5]: @@ -71,13 +64,9 @@ def get_nist_transitions(filename): lambda_angstroms = float(row[1]) else: continue - if len(row[3].strip()) > 0: - A = float(row[3]) - else: - # continue - A = 1e8 - lower_energy_ev, upper_energy_ev = [float(x.strip(" []")) for x in row[5].split("-")] - lower_g, upper_g = [float(x.strip()) for x in row[12].split("-")] + A = float(row[3]) if len(row[3].strip()) > 0 else 1e8 + lower_energy_ev, upper_energy_ev = (float(x.strip(" []")) for x in row[5].split("-")) + lower_g, upper_g = (float(x.strip()) for x in row[12].split("-")) translist.append( transitiontuple(lambda_angstroms, A, lower_energy_ev, upper_energy_ev, lower_g, upper_g) ) @@ -148,7 +137,7 @@ def make_plot( if len(axes) > len(ionlist): axes[len(ionlist)].plot(xvalues, yvalues_combined[seriesindex], linewidth=1.5, label=serieslabel) - peak_y_value = max(peak_y_value, max(yvalues_combined[seriesindex])) + peak_y_value = max(peak_y_value, **yvalues_combined[seriesindex]) axislabels = [ f"{at.get_elsymbol(Z)} {at.roman_numerals[ion_stage]}\n(pop={ionpopdict[(Z, ion_stage)]:.1e}/cm3)" @@ -185,12 +174,13 @@ def make_plot( plt.close() -def add_upper_lte_pop(dftransitions, T_exc, ionpop, ltepartfunc, columnname=None): +def add_upper_lte_pop(dftransitions, T_exc, ionpop, ltepartfunc, columnname=None) -> pd.DataFrame: K_B = const.k_B.to("eV / K").value scalefactor = ionpop / ltepartfunc if columnname is None: columnname = f"upper_pop_lte_{T_exc:.0f}K" - dftransitions.eval(f"{columnname} = @scalefactor * upper_g * exp(-upper_energy_ev / @K_B / @T_exc)", inplace=True) + dftransitions = dftransitions.eval(f"{columnname} = @scalefactor * upper_g * exp(-upper_energy_ev / @K_B / @T_exc)") + return dftransitions def addargs(parser: argparse.ArgumentParser) -> None: @@ -257,20 +247,18 @@ def main(args=None, argsraw=None, **kwargs): if from_model: modelgridindex = args.modelgridindex - if args.timedays: - timestep = at.get_timestep_of_timedays(modelpath, args.timedays) - else: - timestep = args.timestep + timestep = at.get_timestep_of_timedays(modelpath, args.timedays) if args.timedays else args.timestep modeldata, _ = at.inputmodel.get_modeldata(Path(modelpath, "model.txt")) estimators_all = at.estimators.read_estimators(modelpath, timestep=timestep, modelgridindex=modelgridindex) if not estimators_all: - return -1 + print("no estimators") + sys.exit(1) estimators = estimators_all[(timestep, modelgridindex)] if estimators["emptycell"]: print(f"ERROR: cell {modelgridindex} is marked as empty") - return -1 + sys.exit(1) # also calculate wavelengths outside the plot range to include lines whose # edges pass through the plot range @@ -314,7 +302,7 @@ def main(args=None, argsraw=None, **kwargs): if dfnltepops is None or dfnltepops.empty: print(f"ERROR: no NLTE populations for cell {modelgridindex} at timestep {timestep}") - return -1 + sys.exit(1) ionpopdict = { (Z, ion_stage): dfnltepops.query("Z==@Z and ion_stage==@ion_stage")["n_NLTE"].sum() @@ -341,10 +329,7 @@ def main(args=None, argsraw=None, **kwargs): else: if not args.T: args.T = [2000] - if len(args.T) == 1: - figure_title = f"Te = {args.T[0]:.1f}" - else: - figure_title = None + figure_title = f"Te = {args.T[0]:.1f}" if len(args.T) == 1 else None temperature_list = [] vardict = {} @@ -372,8 +357,8 @@ def main(args=None, argsraw=None, **kwargs): ionid = (ion.Z, ion.ion_stage) if ionid not in ionlist: continue - else: - ionindex = ionlist.index(ionid) + + ionindex = ionlist.index(ionid) if args.atomicdatabase == "kurucz": dftransitions = dftransgfall.query("Z == @ion.Z and ionstage == @ion.ion_stage", inplace=False).copy() @@ -388,33 +373,35 @@ def main(args=None, argsraw=None, **kwargs): ) if not args.include_permitted and not dftransitions.empty: - dftransitions.query("forbidden == True", inplace=True) + dftransitions = dftransitions.query("forbidden == True") print(f" ({len(ion.transitions):6d} forbidden)") if not dftransitions.empty: if args.atomicdatabase == "artis": - dftransitions.eval("upper_energy_ev = @ion.levels.loc[upper].energy_ev.values", inplace=True) - dftransitions.eval("lower_energy_ev = @ion.levels.loc[lower].energy_ev.values", inplace=True) - dftransitions.eval("lambda_angstroms = @hc / (upper_energy_ev - lower_energy_ev)", inplace=True) + dftransitions = dftransitions.eval("upper_energy_ev = @ion.levels.loc[upper].energy_ev.to_numpy()") + dftransitions = dftransitions.eval("lower_energy_ev = @ion.levels.loc[lower].energy_ev.to_numpy()") + dftransitions = dftransitions.eval("lambda_angstroms = @hc / (upper_energy_ev - lower_energy_ev)") - dftransitions.query( - "lambda_angstroms >= @plot_xmin_wide & lambda_angstroms <= @plot_xmax_wide", inplace=True + dftransitions = dftransitions.query( + "lambda_angstroms >= @plot_xmin_wide & lambda_angstroms <= @plot_xmax_wide" ) - dftransitions.sort_values(by="lambda_angstroms", inplace=True) + dftransitions = dftransitions.sort_values(by="lambda_angstroms") print(f" {len(dftransitions)} plottable transitions") if args.atomicdatabase == "artis": - dftransitions.eval("upper_g = @ion.levels.loc[upper].g.values", inplace=True) + dftransitions = dftransitions.eval("upper_g = @ion.levels.loc[upper].g.to_numpy()") K_B = const.k_B.to("eV / K").value T_exc = vardict["Te"] ltepartfunc = ion.levels.eval("g * exp(-energy_ev / @K_B / @T_exc)").sum() else: ltepartfunc = 1.0 - dftransitions.eval("flux_factor = (upper_energy_ev - lower_energy_ev) * A", inplace=True) - add_upper_lte_pop(dftransitions, vardict["Te"], ionpopdict[ionid], ltepartfunc, columnname="upper_pop_Te") + dftransitions = dftransitions.eval("flux_factor = (upper_energy_ev - lower_energy_ev) * A") + dftransitions = add_upper_lte_pop( + dftransitions, vardict["Te"], ionpopdict[ionid], ltepartfunc, columnname="upper_pop_Te" + ) for seriesindex, temperature in enumerate(temperature_list): T_exc = eval(temperature, vardict) @@ -424,15 +411,16 @@ def main(args=None, argsraw=None, **kwargs): nltepopdict = {x.level: x["n_NLTE"] for _, x in dfnltepops_thision.iterrows()} dftransitions["upper_pop_nlte"] = dftransitions.apply( - lambda x: nltepopdict.get(x.upper, 0.0), axis=1 + lambda x: nltepopdict.get(x.upper, 0.0), # noqa: B023 # pylint: disable=cell-var-from-loop + axis=1, ) # dftransitions['lower_pop_nlte'] = dftransitions.apply( # lambda x: nltepopdict.get(x.lower, 0.), axis=1) popcolumnname = "upper_pop_nlte" - dftransitions.eval(f"flux_factor_nlte = flux_factor * {popcolumnname}", inplace=True) - dftransitions.eval("upper_departure = upper_pop_nlte / upper_pop_Te", inplace=True) + dftransitions = dftransitions.eval(f"flux_factor_nlte = flux_factor * {popcolumnname}") + dftransitions = dftransitions.eval("upper_departure = upper_pop_nlte / upper_pop_Te") if ionid == (26, 2): fe2depcoeff = dftransitions.query("upper == 16 and lower == 5").iloc[0].upper_departure elif ionid == (28, 2): @@ -443,15 +431,17 @@ def main(args=None, argsraw=None, **kwargs): else: popcolumnname = f"upper_pop_lte_{T_exc:.0f}K" if args.atomicdatabase == "artis": - dftransitions.eval("upper_g = @ion.levels.loc[upper].g.values", inplace=True) + dftransitions = dftransitions.eval("upper_g = @ion.levels.loc[upper].g.to_numpy()") K_B = const.k_B.to("eV / K").value ltepartfunc = ion.levels.eval("g * exp(-energy_ev / @K_B / @T_exc)").sum() else: ltepartfunc = 1.0 - add_upper_lte_pop(dftransitions, T_exc, ionpopdict[ionid], ltepartfunc, columnname=popcolumnname) + dftransitions = add_upper_lte_pop( + dftransitions, T_exc, ionpopdict[ionid], ltepartfunc, columnname=popcolumnname + ) if args.print_lines: - dftransitions.eval(f"flux_factor_{popcolumnname} = flux_factor * {popcolumnname}", inplace=True) + dftransitions = dftransitions.eval(f"flux_factor_{popcolumnname} = flux_factor * {popcolumnname}") yvalues[seriesindex][ionindex] = generate_ion_spectrum( dftransitions, xvalues, popcolumnname, plot_resolution, args @@ -499,10 +489,11 @@ def get_strionfracs(atomic_number, ionstages): f"{estimators['populations'][(28, 3)] / estimators['populations'][(28, 2)]:5.2f}" ) - if from_model: - outputfilename = str(args.outputfile).format(cell=modelgridindex, timestep=timestep, time_days=time_days) - else: - outputfilename = "plottransitions.pdf" + outputfilename = ( + str(args.outputfile).format(cell=modelgridindex, timestep=timestep, time_days=time_days) + if from_model + else "plottransitions.pdf" + ) make_plot( xvalues, diff --git a/artistools/viewing_angles_visualization.py b/artistools/viewing_angles_visualization.py index 0fe58dbe0..4b07665e4 100755 --- a/artistools/viewing_angles_visualization.py +++ b/artistools/viewing_angles_visualization.py @@ -7,7 +7,7 @@ def load_artis_model(modelfile): - with open(modelfile, "r") as f: + with open(modelfile) as f: n = f.readline() res = int(np.cbrt(int(n))) rho = np.zeros([res, res, res]) @@ -151,8 +151,8 @@ def main( isomin, isomax : float, float """ try: - import plotly.graph_objects as go import plotly.express as px + import plotly.graph_objects as go except ModuleNotFoundError: print("Cannot run visualization without plotly...") sys.exit() @@ -180,15 +180,15 @@ def main( animation_frame="Angle-bin", hover_name="Angle-bin", ) - fig.update_traces(line=dict(width=linewidth)) + fig.update_traces(line={"width": linewidth}) fig.update_layout( - legend=dict( - orientation="h", - yanchor="bottom", - y=1.02, - xanchor="right", - x=1, - ) + legend={ + "orientation": "h", + "yanchor": "bottom", + "y": 1.02, + "xanchor": "right", + "x": 1, + } ) fig = fig.add_trace( diff --git a/artistools/writecomparisondata.py b/artistools/writecomparisondata.py index b8a96226b..6726a1bed 100755 --- a/artistools/writecomparisondata.py +++ b/artistools/writecomparisondata.py @@ -3,11 +3,8 @@ from pathlib import Path import numpy as np -import pandas as pd import artistools as at -import artistools.estimators -import artistools.lightcurve def write_spectra(modelpath, model_id, selected_timesteps, outfile): @@ -33,15 +30,15 @@ def write_spectra(modelpath, model_id, selected_timesteps, outfile): lum_lambda[n, :] = fluxes_nu[n, :] * 2.99792458e18 / lambdas[n] / lambdas[n] * area with open(outfile, "w") as f: - f.write("#NTIMES: {0}\n".format(len(selected_timesteps))) - f.write("#NWAVE: {0}\n".format(len(lambdas))) - f.write("#TIMES[d]: {0}\n".format(" ".join(["{0:.2f}".format(times[ts]) for ts in selected_timesteps]))) + f.write("#NTIMES: {}\n".format(len(selected_timesteps))) + f.write("#NWAVE: {}\n".format(len(lambdas))) + f.write("#TIMES[d]: {}\n".format(" ".join(["{:.2f}".format(times[ts]) for ts in selected_timesteps]))) f.write("#wavelength[Ang] flux_t0[erg/s/Ang] flux_t1[erg/s/Ang] ... flux_tn[erg/s/Ang]\n") for n in reversed(range(len(lambdas))): f.write( - "{0:.2f} ".format(lambdas[n]) - + " ".join(["{0:.2e}".format(lum_lambda[n, ts]) for ts in selected_timesteps]) + "{:.2f} ".format(lambdas[n]) + + " ".join(["{:.2e}".format(lum_lambda[n, ts]) for ts in selected_timesteps]) + "\n" ) @@ -156,10 +153,10 @@ def write_phys(modelpath, model_id, selected_timesteps, estimators, allnonemptym def write_lbol_edep(modelpath, model_id, selected_timesteps, estimators, outputpath): # times = at.get_timestep_times_float(modelpath) - dflightcurve = at.lightcurve.readfile(Path(modelpath, "light_curve.out")) + dflightcurve = at.lightcurve.readfile(Path(modelpath, "light_curve.out"))[-1] dfdep = at.get_deposition(modelpath) - df = pd.merge(dflightcurve, dfdep, left_index=True, right_index=True, suffixes=("", "_dep")) + df = dflightcurve.merge(dfdep, left_index=True, right_index=True, suffixes=("", "_dep")) with open(outputpath, "w") as f: f.write(f"#NTIMES: {len(selected_timesteps)}\n") diff --git a/pyproject.toml b/pyproject.toml index 03207d82e..28708bbc1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,15 @@ where = ["."] [tool.black] line-length = 120 -preview = false +preview = true + +[tool.isort] +profile = "black" +src_paths = ["artistools"] +line_length = 120 +force_single_line = true +case_sensitive = false +sort_relative_in_force_sorted_sections = true [tool.mypy] python_version = '3.11' @@ -46,25 +54,113 @@ ignore_missing_imports = true warn_unused_configs = true #files = 'artistools/**/*.py' plugins = 'numpy.typing.mypy_plugin' +scripts_are_modules = true +strict_equality = true +pretty = true +error_summary = true +enable_error_code = [ + "redundant-expr", + "truthy-bool", + "ignore-without-code", +] [[tool.mypy.overrides]] module = [ "artistools.estimators.estimators", "artistools.inputmodel", + "artistools.lightcurve", "artistools.misc", +# "artistools.packets", "artistools.spectra.spectra", ] check_untyped_defs = true disallow_untyped_defs = true warn_return_any = false +#strict = true -[tool.pylint] +[tool.pylint.'MESSAGES CONTROL'] max-line-length = 120 -errors-only = true - +disable = """ + broad-exception-caught, + dangerous-default-value, + eval-used, + fixme, + global-statement, + missing-function-docstring, + missing-module-docstring, + import-outside-toplevel, + invalid-name, + line-too-long, + protected-access, + redefined-builtin, + redefined-outer-name, + too-many-arguments, + too-many-branches, + too-many-lines, + too-many-locals, + too-many-statements, + unused-argument, + unused-variable, + unspecified-encoding, + C, + R, +""" [tool.pylint.typecheck] ignored-modules = ["astropy", "extinction"] +[tool.ruff] +select = [ + #"ALL", + "A", "B", "C", "E", "F", "G", "N", "Q", "W", "ARG", "BLE", "DJ", "DTZ", "EM", "ERA", "EXE", + "ICN", "INP", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PT", "PYI", "RET", "RSE", "RUF", + "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT" +] +ignore = [ + "ARG001", # ignored because variables in df.eval() are not detected + "B005", + "B007", # variable not used in loop body (but sometimes it is with eval) + "BLE001", + "C9", + "D211", "D213", + "E501", "E741", + "EM101","EM102", + "ERA001", + "F841", # ignored because variables in df.eval() are not detected + "N802", "N803", "N806", "N999", + "PD901", + "PGH001", + "PLR0911", "PLR0912", "PLR0913", "PLR0915", "PLR2004", + "PLW2901", + "PT011", + "RET504", + "SIM115", + "SLF001", + "TRY003", + "TRY301", + "UP032", +] + +# Allow autofix for all enabled rules (when `--fix`) is provided. +fixable = [ + "A", "B", "C", "E", "G", "N", "Q", "W", "ARG", "BLE", "DJ", "DTZ", "EM", "ERA", "EXE", + "FBT", "ICN", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PLR", "PT", "PTH", "PYI", "RET", "RSE", + "RUF", "SIM", "SLF", "TCH", "TID", "UP", "YTT" +] +unfixable = [ + "ERA001", # commented out code (will just delete it!) + "F841", + "F401", + "PD002" +] + +# larger than black's limit. Let black handle this +line-length = 200 + +target-version = "py39" + +[tool.ruff.per-file-ignores] +"__init__.py" = ["F401"] + [tool.setuptools_scm] write_to = "_version.py" local_scheme = "no-local-version" diff --git a/requirements.txt b/requirements.txt index 0f6ef7e7a..20977d79d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,25 +1,28 @@ -argcomplete>=2.0.0 -astropy>=5.1 -black>=23.1 +argcomplete>=3.0 +astropy>=5.2 +black>=23.3 coveralls>=3.3.1 extinction>=0.4.6 -flake8>=5.0.4 -matplotlib>=3.5.3 -mypy>=0.971 -numpy>=1.23.2 -pandas>=1.4.3 -pre-commit>=2.20.0 +lz4>=4.3 +matplotlib>=3.7 +mypy>=1.2 +numpy>=1.24 +pandas>=2.0 +polars>=0.16.18 +pre-commit>=3.2 psutil>=5.9.1 -#pyarrow>=9.0.0 +pyarrow>=11.0.0 pynonthermal>=2021.10.12 pypdf2>=2.10.4 -pytest>=7.1.2 -pytest-cov>=3.0.0 +pytest>=7.2 +pytest-cov>=4.0.0 pytest-runner>=6.0.0 -pyvista>=0.36.0 +python-xz>=0.5 +pyvista>=0.36 PyYAML>=6.0 -scipy>=1.9.1 -setuptools_scm[toml]>=7.0.5 -types-PyYAML>=6.0.11 -wheel>=0.37.1 -xattr>=0.9.9 +pyzstd>=0.15 +ruff>=0.0.260 +scipy>=1.10 +setuptools_scm[toml]>=7.1 +types-PyYAML>=6.0 +wheel>=0.40 diff --git a/setup.py b/setup.py index 01e262cf0..5b954e898 100755 --- a/setup.py +++ b/setup.py @@ -2,7 +2,6 @@ # mypy: ignore-errors """Plotting and analysis tools for the ARTIS 3D supernova radiative transfer code.""" import importlib.util -import sys from pathlib import Path from setuptools import find_namespace_packages @@ -22,9 +21,9 @@ packages=find_namespace_packages(where="."), package_dir={"": "."}, url="https://www.github.com/artis-mcrt/artistools/", - long_description=(Path(__file__).absolute().parent / "README.md").open("rt").read(), + long_description=(Path(__file__).absolute().parent / "README.md").open().read(), long_description_content_type="text/markdown", - install_requires=(Path(__file__).absolute().parent / "requirements.txt").open("rt").read().splitlines(), + install_requires=(Path(__file__).absolute().parent / "requirements.txt").open().read().splitlines(), entry_points={ "console_scripts": commands.get_console_scripts(), }, diff --git a/tests/data/.gitignore b/tests/data/.gitignore index d719f35d5..3286bdc76 100644 --- a/tests/data/.gitignore +++ b/tests/data/.gitignore @@ -1,2 +1,3 @@ testmodel -*.tar.* \ No newline at end of file +*.tar.* +*/ \ No newline at end of file diff --git a/tests/data/setuptestdata.sh b/tests/data/setuptestdata.sh index a848b24de..702306b80 100755 --- a/tests/data/setuptestdata.sh +++ b/tests/data/setuptestdata.sh @@ -9,4 +9,7 @@ mkdir -p testmodel/ tar -xf testmodel.tar.xz --directory testmodel/ # find testmodel -size +1M -exec xz -v {} \; +if [ ! -f vspecpolmodel.tar.xz ]; then curl -O https://theory.gsi.de/~lshingle/artis_http_public/artistools/vspecpolmodel.tar.xz; fi +tar -xf vspecpolmodel.tar.xz + set +x \ No newline at end of file