metablate 0.3.0__tar.gz → 0.4.0__tar.gz

{metablate-0.3.0 → metablate-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metablate
-Version: 0.3.0
+Version: 0.4.0
 Summary: A collection of meteor ablation models
 Author-email: Daniel Kastinen <daniel.kastinen@irf.se>, Johan Kero <kero@irf.se>
 Maintainer-email: Daniel Kastinen <daniel.kastinen@irf.se>
@@ -702,6 +702,7 @@ Requires-Dist: xarray
 Requires-Dist: spacecoords>0.2
 Provides-Extra: develop
 Requires-Dist: pytest; extra == "develop"
+Requires-Dist: mypy; extra == "develop"
 Requires-Dist: flake8; extra == "develop"
 Requires-Dist: wheel; extra == "develop"
 Requires-Dist: build; extra == "develop"
@@ -724,6 +725,7 @@ Requires-Dist: mkdocs-section-index; extra == "develop"
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
+Requires-Dist: mypy; extra == "tests"
 Provides-Extra: visualization
 Requires-Dist: pyvista; extra == "visualization"
 Dynamic: license-file

{metablate-0.3.0 → metablate-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "metablate"
-version = "0.3.0"
+version = "0.4.0"
 description = "A collection of meteor ablation models"
 readme = {file = "README.md", content-type = "text/markdown"}
 authors = [
@@ -45,6 +45,7 @@ where = ["src"]
 [project.optional-dependencies]
 develop = [
     "pytest",
+    "mypy",
     "flake8",
     "wheel",
     "build",
@@ -65,7 +66,7 @@ develop = [
     "mkdocs-literate-nav",
     "mkdocs-section-index",
 ]
-tests = ["pytest", "pytest-cov"]
+tests = ["pytest", "pytest-cov", "mypy"]
 visualization = ["pyvista"]
 
 

{metablate-0.3.0 → metablate-0.4.0}/src/metablate/__init__.py

@@ -1,19 +1,17 @@
 """
 # Meteoroid Ablation Models
 """
+from __future__ import annotations
 
 import logging
 import sys
 
-from . import physics
-from . import models
-from . import material
-from . import atmosphere
+from . import atmosphere, material, models, physics
 from .atmosphere import Atmosphere
 from .version import __version__
 
 
-def setup_logging(level=logging.INFO):
+def setup_logging(level: int | str = logging.INFO) -> logging.Logger:
     logger = logging.getLogger("metablate")
     logger.setLevel(level)
 

{metablate-0.3.0 → metablate-0.4.0}/src/metablate/atmosphere/atm_msise00.py

@@ -1,37 +1,44 @@
-#!/usr/bin/env python
-
 """msis00
 
 """
+from __future__ import annotations
+
+from importlib import import_module
+from types import ModuleType
+from typing import Any
+
+import numpy as np
 import scipy.constants
 import xarray as xr
-from .atmosphere import Atmosphere
-from ..types import NDArray_N
 
+from ..types import ArrayLike
+from .atmosphere import Atmosphere
 
+msise00: ModuleType | None
 try:
-    import msise00
+    msise00 = import_module("msise00")
 except ImportError:
     msise00 = None
 
 
 class AtmMSISE00(Atmosphere):
-    def __init__(self):
+
+    def __init__(self) -> None:
         super().__init__(
             supported_species=["He", "O", "N2", "O2", "Ar", "H", "N"],
         )
 
     def density(
         self,
-        time: NDArray_N,
-        lat: NDArray_N,
-        lon: NDArray_N,
-        alt: NDArray_N,
+        time: ArrayLike,
+        lat: ArrayLike,
+        lon: ArrayLike,
+        alt: ArrayLike,
         f107: float = 80.0,
         f107s: float = 80.0,
         Ap: float = 4.0,
         mass_densities: bool = False,
-        **kwargs
+        **kwargs: Any,
     ) -> xr.Dataset:
         """TODO: Write docstring
 
@@ -43,7 +50,7 @@ class AtmMSISE00(Atmosphere):
 
         result = msise00.run(
             time=time,
-            altkm=alt * 1e-3,
+            altkm=np.asarray(alt) * 1e-3,
             glat=lat,
             glon=lon,
             indices=dict(
@@ -51,7 +58,7 @@ class AtmMSISE00(Atmosphere):
                 f107s=f107s,
                 Ap=Ap,
             ),
-            **kwargs
+            **kwargs,
         )
         result["alt_km"] = result["alt_km"]*1e3
         result = result.rename({

{metablate-0.3.0 → metablate-0.4.0}/src/metablate/atmosphere/atm_pymsis.py

@@ -1,15 +1,19 @@
-#!/usr/bin/env python
-
 """pymsis
 
 """
-from typing import Union
-import scipy.constants
-from .atmosphere import Atmosphere
+from __future__ import annotations
+
+from types import ModuleType
+from typing import Any
+
 import numpy as np
-import numpy.typing as npt
+import scipy.constants
 import xarray as xr
 
+from ..types import ArrayLike, AtmosphereData
+from .atmosphere import Atmosphere
+
+msis: ModuleType | None
 try:
     from pymsis import msis
 except ImportError:
@@ -28,25 +32,25 @@ class AtmPymsis(Atmosphere):
     Anomalous oxygen # density (m-3), | NO # density (m-3), 
     | Temperature (K)]
     """
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             supported_species=["He", "O", "N2", "O2", "Ar", "H", "N"],
         )
 
     def density(
         self,
-        time: npt.ArrayLike,
-        lat: npt.ArrayLike,
-        lon: npt.ArrayLike,
-        alt: npt.ArrayLike,
-        f107: npt.ArrayLike | None = None,
-        f107s: npt.ArrayLike | None = None,
-        aps: npt.ArrayLike | None = None,
+        time: ArrayLike,
+        lat: ArrayLike,
+        lon: ArrayLike,
+        alt: ArrayLike,
+        f107: ArrayLike | None = None,
+        f107s: ArrayLike | None = None,
+        aps: ArrayLike | None = None,
         mass_densities: bool = False,
-        version: Union[float, str] = 2.1,
+        version: float | str = 2.1,
         as_dataset: bool = True,
-        **kwargs
-    ) -> xr.Dataset | dict[str, xr.Variable]:
+        **kwargs: Any,
+    ) -> AtmosphereData:
         """Return species number densities [m^-3], or mass densities [kg/m^3].
 
         Set ``as_dataset=False`` in an ablation model's ``atmosphere_kwargs``
@@ -66,7 +70,7 @@ class AtmPymsis(Atmosphere):
             if len(inp.shape) == 0:
                 inp.shape = (1,)
 
-        result = msis.run(
+        data = msis.run(
             dates=time,
             lons=lon,
             lats=lat,
@@ -77,22 +81,23 @@ class AtmPymsis(Atmosphere):
             version=version,
             **kwargs,
         )
-        if len(result.shape) == 2:
-            result = result.reshape(result.shape[0], 1, 1, 1, result.shape[1])
+        if len(data.shape) == 2:
+            data = data.reshape(data.shape[0], 1, 1, 1, data.shape[1])
 
         variables = {
-            "Total": (["time", "lon", "lat", "alt"], result[:, :, :, :, 0]),
-            "N2": (["time", "lon", "lat", "alt"], result[:, :, :, :, 1]),
-            "O2": (["time", "lon", "lat", "alt"], result[:, :, :, :, 2]),
-            "O": (["time", "lon", "lat", "alt"], result[:, :, :, :, 3]),
-            "He": (["time", "lon", "lat", "alt"], result[:, :, :, :, 4]),
-            "H": (["time", "lon", "lat", "alt"], result[:, :, :, :, 5]),
-            "Ar": (["time", "lon", "lat", "alt"], result[:, :, :, :, 6]),
-            "N": (["time", "lon", "lat", "alt"], result[:, :, :, :, 7]),
-            "AnomalousO": (["time", "lon", "lat", "alt"], result[:, :, :, :, 8]),
-            "NO": (["time", "lon", "lat", "alt"], result[:, :, :, :, 9]),
-            "Temperature": (["time", "lon", "lat", "alt"], result[:, :, :, :, 10]),
+            "Total": (["time", "lon", "lat", "alt"], data[:, :, :, :, 0]),
+            "N2": (["time", "lon", "lat", "alt"], data[:, :, :, :, 1]),
+            "O2": (["time", "lon", "lat", "alt"], data[:, :, :, :, 2]),
+            "O": (["time", "lon", "lat", "alt"], data[:, :, :, :, 3]),
+            "He": (["time", "lon", "lat", "alt"], data[:, :, :, :, 4]),
+            "H": (["time", "lon", "lat", "alt"], data[:, :, :, :, 5]),
+            "Ar": (["time", "lon", "lat", "alt"], data[:, :, :, :, 6]),
+            "N": (["time", "lon", "lat", "alt"], data[:, :, :, :, 7]),
+            "AnomalousO": (["time", "lon", "lat", "alt"], data[:, :, :, :, 8]),
+            "NO": (["time", "lon", "lat", "alt"], data[:, :, :, :, 9]),
+            "Temperature": (["time", "lon", "lat", "alt"], data[:, :, :, :, 10]),
         }
+        result: AtmosphereData
         if as_dataset:
             result = xr.Dataset(
                 variables,

{metablate-0.3.0 → metablate-0.4.0}/src/metablate/atmosphere/atmosphere.py

@@ -1,12 +1,15 @@
-from abc import abstractmethod
-from dataclasses import dataclass
+from __future__ import annotations
+
 import logging
+from abc import abstractmethod
 from copy import copy
+from dataclasses import dataclass
+from typing import Any
 
-import scipy.constants as constants
 import numpy as np
-import xarray as xr
-from ..types import NDArray_N
+from scipy import constants
+
+from ..types import ArrayLike, AtmosphereData
 
 logger = logging.getLogger(__name__)
 
@@ -59,7 +62,8 @@ SPECIES = {
 
 # TODO: get a a hold of other atm models to figure out how to generalize interface better
 class Atmosphere:
-    def __init__(self, supported_species: list[Species] | list[str]):
+
+    def __init__(self, supported_species: list[Species] | list[str]) -> None:
         self.species = {}
         for s in supported_species:
             if isinstance(s, Species):
@@ -71,11 +75,6 @@ class Atmosphere:
 
     @abstractmethod
     def density(
-        self,
-        time: np.datetime64,
-        lat: NDArray_N,
-        lon: NDArray_N,
-        alt: NDArray_N,
-        **kwargs
-    ) -> xr.Dataset:
+        self, time: ArrayLike, lat: ArrayLike, lon: ArrayLike, alt: ArrayLike, **kwargs: Any
+    ) -> AtmosphereData:
         pass

{metablate-0.3.0 → metablate-0.4.0}/src/metablate/material.py

@@ -10,16 +10,24 @@ liberated from the meteoroid by thermal ablation were released
 in complete molecular form.
 
 """
+from __future__ import annotations
 
-from dataclasses import asdict
 import logging
-from copy import copy
+from copy import deepcopy
+from dataclasses import asdict
+from typing import Literal, overload
+
 from scipy import constants
-from .types import Material
+
+from .types import Material, MaterialDict
 
 logger = logging.getLogger(__name__)
 
-MATERIALS = {}
+MATERIALS: dict[str, Material] = {}
+
+# Jones (1997), Table 1: abundances by atom number. As in the original
+# ionization model
+JONES_COMPOSITION = {"O": 0.617, "Fe": 0.059, "Mg": 0.082, "Si": 0.242}
 
 iron = Material(
     bulk_density=7800.0,
@@ -32,6 +40,7 @@ iron = Material(
     penetration_correction_parameter_K=0.35,
     clausius_clapeyron_coeff_A=10.6,
     clausius_clapeyron_coeff_B=16120.0,
+    atomic_composition={"Fe": 1.0},
 )
 MATERIALS["iron"] = iron
 
@@ -46,6 +55,7 @@ asteroidal = Material(
     penetration_correction_parameter_K=0.1,
     clausius_clapeyron_coeff_A=10.6,
     clausius_clapeyron_coeff_B=13500.0,
+    atomic_composition={}, # TODO: what should this be?
 )
 MATERIALS["asteroidal"] = asteroidal
 
@@ -60,6 +70,7 @@ cometary = Material(
     penetration_correction_parameter_K=0.65,
     clausius_clapeyron_coeff_A=10.6,
     clausius_clapeyron_coeff_B=13500.0,
+    atomic_composition=JONES_COMPOSITION.copy(),
 )
 MATERIALS["cometary"] = cometary
 
@@ -74,16 +85,29 @@ porous = Material(
     penetration_correction_parameter_K=0.65,
     clausius_clapeyron_coeff_A=10.6,
     clausius_clapeyron_coeff_B=13500.0,
+    atomic_composition=JONES_COMPOSITION.copy(),
 )
 MATERIALS["porous"] = porous
 
 
-def available():
+def available() -> list[str]:
     """List currently available meteoroid materials."""
     return list(MATERIALS.keys())
 
 
-def get(material, as_dict=True):
+@overload
+def get(material: str, as_dict: Literal[True] = True) -> MaterialDict: ...
+
+
+@overload
+def get(material: str, as_dict: Literal[False]) -> Material: ...
+
+
+@overload
+def get(material: str, as_dict: bool) -> Material | MaterialDict: ...
+
+
+def get(material: str, as_dict: bool = True) -> Material | MaterialDict:
     """Returns the physical parameters of the meteoroid based on its material.
 
     Parameters
@@ -96,7 +120,7 @@ def get(material, as_dict=True):
     Returns
     -------
     Material or dict
-        Dictionary of all available material parameters (float or None)
+        Material properties, including the optional atomic composition.
     """
 
     _material = material.lower().strip()
@@ -107,4 +131,4 @@ def get(material, as_dict=True):
     if as_dict:
         return asdict(MATERIALS[_material])
     else:
-        return copy(MATERIALS[_material])
+        return deepcopy(MATERIALS[_material])

{metablate-0.3.0 → metablate-0.4.0}/src/metablate/models/alpha_beta_2026.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import logging
 import time
 from dataclasses import dataclass, field
-from typing import Any
 
 import numpy as np
 import scipy.interpolate as sci
@@ -16,10 +15,15 @@ from scipy.integrate import solve_ivp
 from .. import physics
 from ..atmosphere import Atmosphere
 from ..types import (
+    Bounds,
+    Kwargs,
+    NDArray,
     NDArray_2xN,
     NDArray_N,
+    Numeric,
     Options,
     Results,
+    StateVector,
 )
 from .model import MeteorModel
 
@@ -34,9 +38,9 @@ class AlphaBetaOptions(Options):
     max_time: float = 100.0  # [s]
     target_normalized_velocity: float | None = None  # [1]
     method: str = "RK45"
-    method_kwargs: dict[str, Any] = field(default_factory=lambda: dict())
+    method_kwargs: Kwargs = field(default_factory=dict)
     atmosphere: Atmosphere | None = None
-    atmosphere_kwargs: dict[str, Any] = field(default_factory=lambda: dict())
+    atmosphere_kwargs: Kwargs = field(default_factory=dict)
     atmospheric_scale_height: float = 7610.0
 
 
@@ -62,24 +66,24 @@ class AlphaBetaResults(Results):
 
 
 def dmdt(
-    v: NDArray_N | float,
-    a: NDArray_N | float,
-    beta: NDArray_N | float,
-    mu: NDArray_N | float,
-) -> NDArray_N | float:
+    v: Numeric,
+    a: Numeric,
+    beta: Numeric,
+    mu: Numeric,
+) -> Numeric:
     return 2 * beta * v * a / (1 - mu) * np.exp(-beta * (1 - v**2) / (1 - mu))
 
 
 def diff_eq_rhs(
     t: float,
-    params: list[NDArray_N | float],
+    params: StateVector,
     alpha: float,
     beta: float,
     gamma: float,
     entry_velocity: float,
     atmospheric_scale_height: float,
 ) -> NDArray_N | NDArray_2xN:
-    s, v = params
+    _, v = params
     delta = scs.expi(beta) - scs.expi(beta * v**2)
     coef = entry_velocity / atmospheric_scale_height
     return np.array(
@@ -91,7 +95,8 @@ def diff_eq_rhs(
 
 
 class AlphaBeta2026(MeteorModel[AlphaBetaInitialState, AlphaBetaOptions, AlphaBetaResults]):
-    def __init__(self, options: AlphaBetaOptions = AlphaBetaOptions()):
+
+    def __init__(self, options: AlphaBetaOptions = AlphaBetaOptions()) -> None:
         self.options = options
         if options.atmosphere is not None:
             raise NotImplementedError(
@@ -133,12 +138,12 @@ class AlphaBeta2026(MeteorModel[AlphaBetaInitialState, AlphaBetaOptions, AlphaBe
         s0 = 0
         y0 = np.array([s0, v0_norm])
 
-        def _low_velocity(t: float, y: NDArray_N, *args: tuple[Any]) -> float:
+        def _low_velocity(t: float, y: NDArray_N, *args: object) -> float:
             res = y[1] / v_final_norm - 1
             return res
 
-        _low_velocity.terminal = True  # type: ignore
-        _low_velocity.direction = -1  # type: ignore
+        _low_velocity.terminal = True  # type: ignore[attr-defined]
+        _low_velocity.direction = -1  # type: ignore[attr-defined]
 
         events = [_low_velocity]
 
@@ -179,40 +184,46 @@ class AlphaBeta2026(MeteorModel[AlphaBetaInitialState, AlphaBetaOptions, AlphaBe
             t=ivp_result.t,
             distance=s_vals,
             velocity=v_norm * parameters.entry_velocity,
-            height=physics.alpha_beta.height_direct(
-                v_norm,
-                self.options.atmospheric_scale_height,
-                parameters.alpha,
-                parameters.beta,
+            height=np.asarray(
+                physics.alpha_beta.height_direct(
+                    v_norm,
+                    self.options.atmospheric_scale_height,
+                    parameters.alpha,
+                    parameters.beta,
+                )
             ),
-            relative_mass=physics.alpha_beta.norm_mass_direct(
-                v_norm,
-                parameters.beta,
-                parameters.shape_change_coefficient,
+            relative_mass=np.asarray(
+                physics.alpha_beta.norm_mass_direct(
+                    v_norm,
+                    parameters.beta,
+                    parameters.shape_change_coefficient,
+                )
             ),
-            massloss=dmdt(
-                v_norm,
-                diffs[1, :],
-                parameters.beta,
-                parameters.shape_change_coefficient,
+            massloss=np.asarray(
+                dmdt(
+                    v_norm,
+                    diffs[1, :],
+                    parameters.beta,
+                    parameters.shape_change_coefficient,
+                )
             ),
         )
 
 
 def logposterior_alpha_beta(
-    alpha,
-    beta,
-    heights,
-    heights_std,
-    velocities,
-    velocities_std,
-    t_eval,
-    initial_height,
-    entry_velocity,
-    entry_angle,
-    shape_change_coefficient,
-    model,
-):
+    alpha: float,
+    beta: float,
+    heights: NDArray_N,
+    heights_std: Numeric,
+    velocities: NDArray_N,
+    velocities_std: Numeric,
+    t_eval: NDArray_N,
+    initial_height: float,
+    entry_velocity: float,
+    entry_angle: float,
+    shape_change_coefficient: float,
+    model: AlphaBeta2026,
+) -> float:
     p = AlphaBetaInitialState(
         epoch=None,
         alpha=alpha,
@@ -243,19 +254,19 @@ def logposterior_alpha_beta(
 
 
 def solve_alpha_beta_posterior(
-    t,
-    heights,
-    heights_std,
-    velocities,
-    velocities_std,
-    entry_angle,
-    shape_change_coefficient,
-    entry_velocity=None,
-    options=None,
-    start=None,
-    bounds=((0.01, 1e8), (0.5e-1, 200.0), (0.0, None)),
-    minimize_kwargs={},
-):
+    t: NDArray_N,
+    heights: NDArray_N,
+    heights_std: Numeric,
+    velocities: NDArray_N,
+    velocities_std: Numeric,
+    entry_angle: float,
+    shape_change_coefficient: float,
+    entry_velocity: float | None = None,
+    options: AlphaBetaOptions | None = None,
+    start: NDArray | list[float | None] | None = None,
+    bounds: Bounds = ((0.01, 1e8), (0.5e-1, 200.0), (0.0, None)),
+    minimize_kwargs: Kwargs = {},
+) -> sco.OptimizeResult:
     """Solve for alpha and beta (and optionally velocity) using minimization of the posterior [^1].
 
     [^1]: Kastinen et al (2026) in preparation
@@ -272,10 +283,12 @@ def solve_alpha_beta_posterior(
     if start is None:
         start = [a0, b0, v0]
     else:
-        if start[1] is None:
-            start[1] = b0
+        beta_start = start[1]
+        if beta_start is None:
+            beta_start = b0
+            start[1] = beta_start
         if start[0] is None:
-            start[0] = np.exp(Yvalues[-1]) / (2.0 * start[1])
+            start[0] = np.exp(Yvalues[-1]) / (2.0 * beta_start)
         if start[2] is None:
             start[2] = v0
 
@@ -286,7 +299,7 @@ def solve_alpha_beta_posterior(
         start = start[:2]
         bounds = (bounds[0], bounds[1])
 
-        def wrap_func(x):
+        def wrap_func(x: NDArray_N) -> float:
             return -logposterior_alpha_beta(
                 x[0],
                 x[1],
@@ -301,9 +314,10 @@ def solve_alpha_beta_posterior(
                 shape_change_coefficient,
                 model,
             )
+
     else:
 
-        def wrap_func(x):
+        def wrap_func(x: NDArray_N) -> float:
             return -logposterior_alpha_beta(
                 x[0],
                 x[1],