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
artis-mcrt · Apr 19, 2023 · f32d7cb · f32d7cb
1 parent 56d3edd
commit f32d7cb
Show file tree

Hide file tree

Showing 92 changed files with 4,253 additions and 3,964 deletions.
diff --git 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
@@ -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
@@ -18,6 +18,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 coverage.xml
+.*_cache/
 
 *.so
 
@@ -61,4 +62,5 @@ plottingscripts/
 
 _version.py
 
-.vscode
+.vscode
+.dmypy.json
diff --git 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,38 +28,25 @@ 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:
           - id: reorder-python-imports
             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
@@ -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
@@ -1,9 +1,7 @@
-#!/usr/bin/env python3
 # PYTHON_ARGCOMPLETE_OK
 import argparse
 import importlib
 import multiprocessing
-from typing import Union
 
 import argcomplete
 
@@ -16,37 +14,40 @@ 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"),
     },
     "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))
 
 
 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
@@ -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