metablate 0.1.0__tar.gz → 0.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. {metablate-0.1.0 → metablate-0.4.0}/PKG-INFO +79 -21
  2. metablate-0.4.0/README.md +92 -0
  3. {metablate-0.1.0 → metablate-0.4.0}/pyproject.toml +30 -4
  4. metablate-0.4.0/src/metablate/__init__.py +25 -0
  5. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/atmosphere/atm_msise00.py +22 -13
  6. metablate-0.4.0/src/metablate/atmosphere/atm_pymsis.py +126 -0
  7. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/atmosphere/atmosphere.py +12 -12
  8. metablate-0.4.0/src/metablate/material.py +134 -0
  9. metablate-0.4.0/src/metablate/models/__init__.py +6 -0
  10. metablate-0.4.0/src/metablate/models/alpha_beta_2026.py +337 -0
  11. metablate-0.4.0/src/metablate/models/dimant_oppenheim_2017.py +389 -0
  12. metablate-0.4.0/src/metablate/models/kero_szasz_2008.py +374 -0
  13. metablate-0.4.0/src/metablate/models/model.py +16 -0
  14. metablate-0.4.0/src/metablate/models/simple_ablation_2026.py +303 -0
  15. metablate-0.4.0/src/metablate/models/simple_motion.py +210 -0
  16. metablate-0.4.0/src/metablate/models/stulov_mirskii_vislyi_1995.py +229 -0
  17. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/physics/__init__.py +4 -1
  18. metablate-0.4.0/src/metablate/physics/aerodynamics.py +72 -0
  19. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/physics/alpha_beta.py +258 -215
  20. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/physics/ballistic.py +62 -30
  21. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/physics/luminosity.py +3 -1
  22. metablate-0.4.0/src/metablate/physics/plasma_production.py +281 -0
  23. metablate-0.4.0/src/metablate/physics/spheres.py +13 -0
  24. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/physics/sputtering.py +14 -10
  25. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/physics/thermal_ablation.py +225 -157
  26. metablate-0.4.0/src/metablate/types.py +99 -0
  27. metablate-0.4.0/src/metablate/version.py +6 -0
  28. {metablate-0.1.0 → metablate-0.4.0}/src/metablate.egg-info/PKG-INFO +79 -21
  29. {metablate-0.1.0 → metablate-0.4.0}/src/metablate.egg-info/SOURCES.txt +20 -4
  30. {metablate-0.1.0 → metablate-0.4.0}/src/metablate.egg-info/requires.txt +9 -2
  31. metablate-0.4.0/tests/test_aerodynamics.py +19 -0
  32. metablate-0.4.0/tests/test_alpha_beta.py +85 -0
  33. metablate-0.4.0/tests/test_atmosphere.py +72 -0
  34. metablate-0.4.0/tests/test_ballistic.py +26 -0
  35. metablate-0.4.0/tests/test_dimant_oppenheim.py +83 -0
  36. metablate-0.4.0/tests/test_dimant_oppenheim_paper.py +61 -0
  37. metablate-0.4.0/tests/test_dimant_oppenheim_points.py +31 -0
  38. metablate-0.4.0/tests/test_material.py +17 -0
  39. metablate-0.4.0/tests/test_models.py +99 -0
  40. metablate-0.4.0/tests/test_plasma_production.py +25 -0
  41. metablate-0.4.0/tests/test_sputtering.py +27 -0
  42. metablate-0.4.0/tests/test_thermal_ablation.py +34 -0
  43. metablate-0.1.0/README.md +0 -40
  44. metablate-0.1.0/src/metablate/__init__.py +0 -16
  45. metablate-0.1.0/src/metablate/ablation_model.py +0 -40
  46. metablate-0.1.0/src/metablate/atmosphere/atm_pymsis.py +0 -114
  47. metablate-0.1.0/src/metablate/coordinates.py +0 -149
  48. metablate-0.1.0/src/metablate/material.py +0 -155
  49. metablate-0.1.0/src/metablate/models/__init__.py +0 -5
  50. metablate-0.1.0/src/metablate/models/dimant_oppenheim_2017.py +0 -293
  51. metablate-0.1.0/src/metablate/models/kero_szasz_2008.py +0 -348
  52. metablate-0.1.0/src/metablate/ode.py +0 -82
  53. metablate-0.1.0/src/metablate/version.py +0 -3
  54. metablate-0.1.0/tests/test_atmosphere.py +0 -23
  55. metablate-0.1.0/tests/test_ode_solver.py +0 -27
  56. {metablate-0.1.0 → metablate-0.4.0}/LICENSE +0 -0
  57. {metablate-0.1.0 → metablate-0.4.0}/setup.cfg +0 -0
  58. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/atmosphere/__init__.py +0 -0
  59. {metablate-0.1.0 → metablate-0.4.0}/src/metablate/py.typed +0 -0
  60. {metablate-0.1.0 → metablate-0.4.0}/src/metablate.egg-info/dependency_links.txt +0 -0
  61. {metablate-0.1.0 → metablate-0.4.0}/src/metablate.egg-info/top_level.txt +0 -0
  62. {metablate-0.1.0 → metablate-0.4.0}/tests/test_import.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: metablate
3
- Version: 0.1.0
3
+ Version: 0.4.0
4
4
  Summary: A collection of meteor ablation models
5
5
  Author-email: Daniel Kastinen <daniel.kastinen@irf.se>, Johan Kero <kero@irf.se>
6
6
  Maintainer-email: Daniel Kastinen <daniel.kastinen@irf.se>
@@ -694,12 +694,15 @@ Description-Content-Type: text/markdown
694
694
  License-File: LICENSE
695
695
  Requires-Dist: tqdm
696
696
  Requires-Dist: scipy
697
- Requires-Dist: numpy
697
+ Requires-Dist: numpy>2
698
+ Requires-Dist: numba>=0.60
698
699
  Requires-Dist: matplotlib
699
- Requires-Dist: xarray
700
700
  Requires-Dist: pymsis
701
+ Requires-Dist: xarray
702
+ Requires-Dist: spacecoords>0.2
701
703
  Provides-Extra: develop
702
704
  Requires-Dist: pytest; extra == "develop"
705
+ Requires-Dist: mypy; extra == "develop"
703
706
  Requires-Dist: flake8; extra == "develop"
704
707
  Requires-Dist: wheel; extra == "develop"
705
708
  Requires-Dist: build; extra == "develop"
@@ -722,45 +725,100 @@ Requires-Dist: mkdocs-section-index; extra == "develop"
722
725
  Provides-Extra: tests
723
726
  Requires-Dist: pytest; extra == "tests"
724
727
  Requires-Dist: pytest-cov; extra == "tests"
728
+ Requires-Dist: mypy; extra == "tests"
729
+ Provides-Extra: visualization
730
+ Requires-Dist: pyvista; extra == "visualization"
725
731
  Dynamic: license-file
726
732
 
727
733
  # Meteoroid Ablation Models
728
734
 
735
+ `metablate` provides meteoroid ablation and atmospheric-entry models, MSIS
736
+ atmospheres, material presets, and plasma-density calculations.
737
+
729
738
  ## Install
730
739
 
740
+ Install the latest published release:
741
+
731
742
  ```bash
732
- pip install ablate
743
+ pip install metablate
733
744
  ```
734
745
 
735
- ### matplotlib backends
746
+ The bleeding-edge version is available with:
736
747
 
737
- on arch to get matplotlib to work
748
+ ```bash
749
+ pip install git+https://github.com/danielk333/ablate.git
750
+ ```
751
+
752
+ or trough cloning this repo.
738
753
 
739
- `pip install Qt5Agg`
754
+ ## Minimal ablation example
740
755
 
741
- and add a `~/.config/matplotlib/matplotlibrc` with `backend:qt5agg`.
756
+ Run the Kero–Szasz (2008) model with a cometary meteoroid and the default
757
+ `AtmPymsis` atmosphere:
742
758
 
743
- see `https://matplotlib.org/stable/tutorials/introductory/usage.html#backends`
759
+ ```python
760
+ import numpy as np
761
+ from spacecoords import frames
744
762
 
745
- ### msise00 atmospheric model
763
+ import metablate.material as mat
764
+ import metablate.models.kero_szasz_2008 as ks
746
765
 
747
- make sure cmake is installed! `sudo pacman -S cmake` and also gfortran `sudo pacman -S gcc gcc-fortran` and the basic libs `sudo pacman -S base-devel`
766
+ lat, lon = 67.0, 20.0
767
+ position = frames.geodetic_wgs84_to_ecef(lat, lon, 100e3, degrees=True)
768
+ direction = frames.azel_to_ecef(lat, lon, az=10, el=-45, degrees=True)
769
+
770
+ model = ks.KeroSzasz2008(
771
+ options=ks.KeroSzaszOptions(
772
+ material=mat.cometary,
773
+ atmosphere_kwargs=dict(as_dataset=False),
774
+ ),
775
+ )
776
+ result = model.run(
777
+ parameters=ks.KeroSzaszInitialState(
778
+ epoch=np.datetime64("2018-06-28T12:45:33"),
779
+ position_ecef=position,
780
+ velocity_ecef=direction * 60e3,
781
+ mass=1e-6,
782
+ ),
783
+ )
784
+ print(f"Runtime: {result.runtime:.3f} s")
785
+ ```
748
786
 
749
- if compilation fails use the enviormennt variables, for bash `export FC="/usr/bin/gfortran"` in fish `set -x FC "/usr/bin/gfortran"`. If cmake still tries to use gcc for compilation, make sure to remove the build folder under `my_env/lib/pythonX.X/site-packages/msise00/build` so that no cached files override the enviornment variables.
787
+ Positions and velocities use ECEF coordinates in metres and metres per second;
788
+ mass is in kilograms. Results include time, mass, velocity, temperature and position.
750
789
 
790
+ ## Atmospheres and computation
791
+
792
+ `AtmPymsis` uses `pymsis`, with MSIS 2.1 by default. Historical solar and geomagnetic
793
+ indices are downloaded and cached when not supplied. Setting `as_dataset=False`
794
+ in `atmosphere_kwargs` avoids constructing an xarray dataset at each solver step.
795
+
796
+ `AtmMSISE00` is optional and requires the separate `msise00` package and a Fortran
797
+ compiler; see its [installation instructions](https://github.com/space-physics/msise00#install).
798
+
799
+ Some functions use Numba to speed up subsequent calls. The first call
800
+ compiles the kernels, which are cached for subsequent runs. Set `NUMBA_DISABLE_JIT=1`
801
+ to run the same formulas in Python for debugging.
802
+
803
+ ## Examples
804
+
805
+ See [docs/examples](docs/examples) for e.g. ablation, fireball, alpha–beta and plasma examples.
806
+ For example, after a repository checkout, generate an ablation sweep with:
751
807
 
752
808
  ```bash
753
- mkdir /my/env/msise00_source
754
- cd /my/env/msise00_source
755
- git clone https://github.com/scivision/msise00
756
- cd msise00
757
- pip install -e .
758
- python -c "import msise00; msise00.base.build()"
809
+ python docs/examples/ablation_sweep.py --output ablation_sweep.png
759
810
  ```
760
811
 
761
- or
812
+ Which produces this plot:
813
+
814
+ ![Kero-Szasz ablation sweep](docs/assets/kero_figure1_example.png)
815
+
816
+ ## Development
817
+
818
+ From a repository checkout:
762
819
 
763
820
  ```bash
764
- pip install git+git://github.com/space-physics/msise00.git@main
765
- python -c "import msise00; msise00.base.build()"
821
+ pip install -e .[tests]
822
+ # or pip install -e .[develop] for all dev packages
823
+ pytest
766
824
  ```
@@ -0,0 +1,92 @@
1
+ # Meteoroid Ablation Models
2
+
3
+ `metablate` provides meteoroid ablation and atmospheric-entry models, MSIS
4
+ atmospheres, material presets, and plasma-density calculations.
5
+
6
+ ## Install
7
+
8
+ Install the latest published release:
9
+
10
+ ```bash
11
+ pip install metablate
12
+ ```
13
+
14
+ The bleeding-edge version is available with:
15
+
16
+ ```bash
17
+ pip install git+https://github.com/danielk333/ablate.git
18
+ ```
19
+
20
+ or trough cloning this repo.
21
+
22
+ ## Minimal ablation example
23
+
24
+ Run the Kero–Szasz (2008) model with a cometary meteoroid and the default
25
+ `AtmPymsis` atmosphere:
26
+
27
+ ```python
28
+ import numpy as np
29
+ from spacecoords import frames
30
+
31
+ import metablate.material as mat
32
+ import metablate.models.kero_szasz_2008 as ks
33
+
34
+ lat, lon = 67.0, 20.0
35
+ position = frames.geodetic_wgs84_to_ecef(lat, lon, 100e3, degrees=True)
36
+ direction = frames.azel_to_ecef(lat, lon, az=10, el=-45, degrees=True)
37
+
38
+ model = ks.KeroSzasz2008(
39
+ options=ks.KeroSzaszOptions(
40
+ material=mat.cometary,
41
+ atmosphere_kwargs=dict(as_dataset=False),
42
+ ),
43
+ )
44
+ result = model.run(
45
+ parameters=ks.KeroSzaszInitialState(
46
+ epoch=np.datetime64("2018-06-28T12:45:33"),
47
+ position_ecef=position,
48
+ velocity_ecef=direction * 60e3,
49
+ mass=1e-6,
50
+ ),
51
+ )
52
+ print(f"Runtime: {result.runtime:.3f} s")
53
+ ```
54
+
55
+ Positions and velocities use ECEF coordinates in metres and metres per second;
56
+ mass is in kilograms. Results include time, mass, velocity, temperature and position.
57
+
58
+ ## Atmospheres and computation
59
+
60
+ `AtmPymsis` uses `pymsis`, with MSIS 2.1 by default. Historical solar and geomagnetic
61
+ indices are downloaded and cached when not supplied. Setting `as_dataset=False`
62
+ in `atmosphere_kwargs` avoids constructing an xarray dataset at each solver step.
63
+
64
+ `AtmMSISE00` is optional and requires the separate `msise00` package and a Fortran
65
+ compiler; see its [installation instructions](https://github.com/space-physics/msise00#install).
66
+
67
+ Some functions use Numba to speed up subsequent calls. The first call
68
+ compiles the kernels, which are cached for subsequent runs. Set `NUMBA_DISABLE_JIT=1`
69
+ to run the same formulas in Python for debugging.
70
+
71
+ ## Examples
72
+
73
+ See [docs/examples](docs/examples) for e.g. ablation, fireball, alpha–beta and plasma examples.
74
+ For example, after a repository checkout, generate an ablation sweep with:
75
+
76
+ ```bash
77
+ python docs/examples/ablation_sweep.py --output ablation_sweep.png
78
+ ```
79
+
80
+ Which produces this plot:
81
+
82
+ ![Kero-Szasz ablation sweep](docs/assets/kero_figure1_example.png)
83
+
84
+ ## Development
85
+
86
+ From a repository checkout:
87
+
88
+ ```bash
89
+ pip install -e .[tests]
90
+ # or pip install -e .[develop] for all dev packages
91
+ pytest
92
+ ```
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "metablate"
7
- version = "0.1.0"
7
+ version = "0.4.0"
8
8
  description = "A collection of meteor ablation models"
9
9
  readme = {file = "README.md", content-type = "text/markdown"}
10
10
  authors = [
@@ -27,10 +27,12 @@ requires-python = ">=3.9"
27
27
  dependencies = [
28
28
  "tqdm",
29
29
  "scipy",
30
- "numpy",
30
+ "numpy>2",
31
+ "numba>=0.60",
31
32
  "matplotlib",
32
- "xarray",
33
33
  "pymsis",
34
+ "xarray",
35
+ "spacecoords>0.2",
34
36
  ]
35
37
 
36
38
  [tool.setuptools]
@@ -43,6 +45,7 @@ where = ["src"]
43
45
  [project.optional-dependencies]
44
46
  develop = [
45
47
  "pytest",
48
+ "mypy",
46
49
  "flake8",
47
50
  "wheel",
48
51
  "build",
@@ -63,7 +66,8 @@ develop = [
63
66
  "mkdocs-literate-nav",
64
67
  "mkdocs-section-index",
65
68
  ]
66
- tests = ["pytest", "pytest-cov"]
69
+ tests = ["pytest", "pytest-cov", "mypy"]
70
+ visualization = ["pyvista"]
67
71
 
68
72
 
69
73
  [project.urls]
@@ -74,3 +78,25 @@ Issues = "https://github.com/danielk333/ablate/issues"
74
78
 
75
79
  [tool.black]
76
80
  line-length = 100
81
+
82
+ [tool.mypy]
83
+ disallow_untyped_defs = true
84
+ follow_untyped_imports = true
85
+ warn_unused_ignores = true
86
+ packages = ["metablate"]
87
+
88
+ [tool.ruff]
89
+ line-length = 110
90
+
91
+ [tool.ruff.lint.per-file-ignores]
92
+ "__init__.py" = ["F401"]
93
+
94
+ [tool.ruff.lint]
95
+ ignore = ["E203","D203","E251","E501","E226"]
96
+
97
+ [tool.pytest]
98
+ addopts = ["-v", "--cov=metablate"]
99
+ testpaths = ["tests"]
100
+ markers = [
101
+ "slow: tests expected to take longer than simple unit checks",
102
+ ]
@@ -0,0 +1,25 @@
1
+ """
2
+ # Meteoroid Ablation Models
3
+ """
4
+ from __future__ import annotations
5
+
6
+ import logging
7
+ import sys
8
+
9
+ from . import atmosphere, material, models, physics
10
+ from .atmosphere import Atmosphere
11
+ from .version import __version__
12
+
13
+
14
+ def setup_logging(level: int | str = logging.INFO) -> logging.Logger:
15
+ logger = logging.getLogger("metablate")
16
+ logger.setLevel(level)
17
+
18
+ formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
19
+
20
+ handler = logging.StreamHandler(sys.stdout)
21
+ handler.setLevel(level)
22
+ handler.setFormatter(formatter)
23
+ logger.addHandler(handler)
24
+
25
+ return logger
@@ -1,36 +1,45 @@
1
- #!/usr/bin/env python
2
-
3
1
  """msis00
4
2
 
5
3
  """
4
+ from __future__ import annotations
5
+
6
+ from importlib import import_module
7
+ from types import ModuleType
8
+ from typing import Any
9
+
10
+ import numpy as np
6
11
  import scipy.constants
12
+ import xarray as xr
13
+
14
+ from ..types import ArrayLike
7
15
  from .atmosphere import Atmosphere
8
- import numpy.typing as npt
9
16
 
17
+ msise00: ModuleType | None
10
18
  try:
11
- import msise00
19
+ msise00 = import_module("msise00")
12
20
  except ImportError:
13
21
  msise00 = None
14
22
 
15
23
 
16
24
  class AtmMSISE00(Atmosphere):
17
- def __init__(self):
25
+
26
+ def __init__(self) -> None:
18
27
  super().__init__(
19
28
  supported_species=["He", "O", "N2", "O2", "Ar", "H", "N"],
20
29
  )
21
30
 
22
31
  def density(
23
32
  self,
24
- time: npt.ArrayLike,
25
- lat: npt.ArrayLike,
26
- lon: npt.ArrayLike,
27
- alt: npt.ArrayLike,
33
+ time: ArrayLike,
34
+ lat: ArrayLike,
35
+ lon: ArrayLike,
36
+ alt: ArrayLike,
28
37
  f107: float = 80.0,
29
38
  f107s: float = 80.0,
30
39
  Ap: float = 4.0,
31
40
  mass_densities: bool = False,
32
- **kwargs
33
- ):
41
+ **kwargs: Any,
42
+ ) -> xr.Dataset:
34
43
  """TODO: Write docstring
35
44
 
36
45
  returns density in [m^-3]
@@ -41,7 +50,7 @@ class AtmMSISE00(Atmosphere):
41
50
 
42
51
  result = msise00.run(
43
52
  time=time,
44
- altkm=alt * 1e-3,
53
+ altkm=np.asarray(alt) * 1e-3,
45
54
  glat=lat,
46
55
  glon=lon,
47
56
  indices=dict(
@@ -49,7 +58,7 @@ class AtmMSISE00(Atmosphere):
49
58
  f107s=f107s,
50
59
  Ap=Ap,
51
60
  ),
52
- **kwargs
61
+ **kwargs,
53
62
  )
54
63
  result["alt_km"] = result["alt_km"]*1e3
55
64
  result = result.rename({
@@ -0,0 +1,126 @@
1
+ """pymsis
2
+
3
+ """
4
+ from __future__ import annotations
5
+
6
+ from types import ModuleType
7
+ from typing import Any
8
+
9
+ import numpy as np
10
+ import scipy.constants
11
+ import xarray as xr
12
+
13
+ from ..types import ArrayLike, AtmosphereData
14
+ from .atmosphere import Atmosphere
15
+
16
+ msis: ModuleType | None
17
+ try:
18
+ from pymsis import msis
19
+ except ImportError:
20
+ msis = None
21
+
22
+
23
+ class AtmPymsis(Atmosphere):
24
+ """
25
+ # TODO: docstring
26
+
27
+ ndarray (ndates, nlons, nlats, nalts, 11) or (ndates, 11) |
28
+ The data calculated at each grid point:
29
+ | [Total mass density (kg/m3) | N2 # density (m-3), | O2 # density (m-3), |
30
+ O # density (m-3), | He # density (m-3), | H # density (m-3), |
31
+ Ar # density (m-3), | N # density (m-3), |
32
+ Anomalous oxygen # density (m-3), | NO # density (m-3),
33
+ | Temperature (K)]
34
+ """
35
+ def __init__(self) -> None:
36
+ super().__init__(
37
+ supported_species=["He", "O", "N2", "O2", "Ar", "H", "N"],
38
+ )
39
+
40
+ def density(
41
+ self,
42
+ time: ArrayLike,
43
+ lat: ArrayLike,
44
+ lon: ArrayLike,
45
+ alt: ArrayLike,
46
+ f107: ArrayLike | None = None,
47
+ f107s: ArrayLike | None = None,
48
+ aps: ArrayLike | None = None,
49
+ mass_densities: bool = False,
50
+ version: float | str = 2.1,
51
+ as_dataset: bool = True,
52
+ **kwargs: Any,
53
+ ) -> AtmosphereData:
54
+ """Return species number densities [m^-3], or mass densities [kg/m^3].
55
+
56
+ Set ``as_dataset=False`` in an ablation model's ``atmosphere_kwargs``
57
+ to skip dataset and coordinate construction at each solver step.
58
+ The returned variables still support ``.values``; total mass density
59
+ and temperature retain their units of kg/m^3 and K, respectively.
60
+ """
61
+ if msis is None:
62
+ raise ImportError("pymsis is not installed")
63
+
64
+ time = np.array(time) if not isinstance(time, np.ndarray) else time
65
+ lon = np.array(lon) if not isinstance(lon, np.ndarray) else lon
66
+ lat = np.array(lat) if not isinstance(lat, np.ndarray) else lat
67
+ alt = np.array(alt) if not isinstance(alt, np.ndarray) else alt
68
+
69
+ for inp in [time, lon, lat, alt]:
70
+ if len(inp.shape) == 0:
71
+ inp.shape = (1,)
72
+
73
+ data = msis.run(
74
+ dates=time,
75
+ lons=lon,
76
+ lats=lat,
77
+ alts=alt * 1e-3,
78
+ f107s=f107,
79
+ f107as=f107s,
80
+ aps=aps,
81
+ version=version,
82
+ **kwargs,
83
+ )
84
+ if len(data.shape) == 2:
85
+ data = data.reshape(data.shape[0], 1, 1, 1, data.shape[1])
86
+
87
+ variables = {
88
+ "Total": (["time", "lon", "lat", "alt"], data[:, :, :, :, 0]),
89
+ "N2": (["time", "lon", "lat", "alt"], data[:, :, :, :, 1]),
90
+ "O2": (["time", "lon", "lat", "alt"], data[:, :, :, :, 2]),
91
+ "O": (["time", "lon", "lat", "alt"], data[:, :, :, :, 3]),
92
+ "He": (["time", "lon", "lat", "alt"], data[:, :, :, :, 4]),
93
+ "H": (["time", "lon", "lat", "alt"], data[:, :, :, :, 5]),
94
+ "Ar": (["time", "lon", "lat", "alt"], data[:, :, :, :, 6]),
95
+ "N": (["time", "lon", "lat", "alt"], data[:, :, :, :, 7]),
96
+ "AnomalousO": (["time", "lon", "lat", "alt"], data[:, :, :, :, 8]),
97
+ "NO": (["time", "lon", "lat", "alt"], data[:, :, :, :, 9]),
98
+ "Temperature": (["time", "lon", "lat", "alt"], data[:, :, :, :, 10]),
99
+ }
100
+ result: AtmosphereData
101
+ if as_dataset:
102
+ result = xr.Dataset(
103
+ variables,
104
+ coords={
105
+ "time": time,
106
+ "lon": lon,
107
+ "lat": lat,
108
+ "alt": alt,
109
+ },
110
+ attrs=dict(
111
+ f107=f107,
112
+ f107s=f107s,
113
+ aps=aps,
114
+ mass_densities=mass_densities,
115
+ version=version,
116
+ ),
117
+ )
118
+ else:
119
+ result = {key: xr.Variable(*value) for key, value in variables.items()}
120
+
121
+ if mass_densities:
122
+ for symbol, s in self.species.items():
123
+ weight = scipy.constants.u * s.A
124
+ result[symbol] = result[symbol] * weight
125
+
126
+ return result
@@ -1,11 +1,15 @@
1
- from abc import abstractmethod
2
- from dataclasses import dataclass
1
+ from __future__ import annotations
2
+
3
3
  import logging
4
+ from abc import abstractmethod
4
5
  from copy import copy
6
+ from dataclasses import dataclass
7
+ from typing import Any
5
8
 
6
- import scipy.constants as constants
7
9
  import numpy as np
8
- import numpy.typing as npt
10
+ from scipy import constants
11
+
12
+ from ..types import ArrayLike, AtmosphereData
9
13
 
10
14
  logger = logging.getLogger(__name__)
11
15
 
@@ -58,7 +62,8 @@ SPECIES = {
58
62
 
59
63
  # TODO: get a a hold of other atm models to figure out how to generalize interface better
60
64
  class Atmosphere:
61
- def __init__(self, supported_species):
65
+
66
+ def __init__(self, supported_species: list[Species] | list[str]) -> None:
62
67
  self.species = {}
63
68
  for s in supported_species:
64
69
  if isinstance(s, Species):
@@ -70,11 +75,6 @@ class Atmosphere:
70
75
 
71
76
  @abstractmethod
72
77
  def density(
73
- self,
74
- time: npt.ArrayLike,
75
- lat: npt.ArrayLike,
76
- lon: npt.ArrayLike,
77
- alt: npt.ArrayLike,
78
- **kwargs
79
- ):
78
+ self, time: ArrayLike, lat: ArrayLike, lon: ArrayLike, alt: ArrayLike, **kwargs: Any
79
+ ) -> AtmosphereData:
80
80
  pass