refactor: modernise type annotations and Streamlit API usage

  - Replace all `Optional[X]` with `X | None` and drop `typing.Optional`
    imports across src/ and tests/
  - Use `astropy.constants` module import instead of direct `m_p` symbol
    import in drag.py
  - Simplify M-is-None guard clauses to ternary expressions
  - Replace `try/except ValueError: pass` with `contextlib.suppress`
  - Update `st.plotly_chart(use_container_width=True)` → `width="stretch"`
    and `st.button/data_editor(use_container_width=True)` → `width="content"`
    to match the current Streamlit API
  - Fix pytest.mark.parametrize argument tuples (string → tuple literal)
    and add `match=` to bare `pytest.raises` calls in tests
  - Split compound assert into separate assertions for clearer failure msgs

Author: ennomuller

pages/01_Propagation_Simulator.py

@@ -7,8 +7,8 @@
 
 from __future__ import annotations
 
+import contextlib
 from datetime import datetime
-from typing import Optional
 
 import astropy.units as u
 import pandas as pd
@@ -119,7 +119,7 @@ def _mini_metric(
 has_obs = len(_obs_clean) >= 2
 
 if has_obs:
-    clean_df: Optional[pd.DataFrame] = _obs_clean.copy()
+    clean_df = _obs_clean.copy()
     clean_df["lon"] = gcs_params.lon_deg
     clean_df["lat"] = gcs_params.lat_deg
     clean_df["tilt"] = gcs_params.tilt_deg
@@ -190,7 +190,7 @@ def _mini_metric(
             target_name=config.target.name,
             height=1.0,
         )
-        st.plotly_chart(gcs_fig, use_container_width=True)
+        st.plotly_chart(gcs_fig, width="stretch")
 
     # --------------------------------------------------------
     # Height-Time Diagram — shown once ≥2 observations exist
@@ -217,7 +217,7 @@ def _mini_metric(
             v_apex_error_kms=derived.v_apex_error_kms,
             event_label=config.event_str,
         )
-        st.plotly_chart(ht_fig, use_container_width=True)
+        st.plotly_chart(ht_fig, width="stretch")
 
         # Velocity result displayed prominently below the plot
         v_col, proj_col, target_col, _ = st.columns([1, 1, 1, 1])
@@ -273,8 +273,8 @@ def _mini_metric(
     # --------------------------------------------------------
     # Display stored results
     # --------------------------------------------------------
-    results: Optional[SimulationResults] = st.session_state.get(KEY_SIM_RESULTS)
-    stored_config: Optional[SimulationConfig] = st.session_state.get(KEY_SIM_CONFIG)
+    results: SimulationResults | None = st.session_state.get(KEY_SIM_RESULTS)
+    stored_config: SimulationConfig | None = st.session_state.get(KEY_SIM_CONFIG)
 
     if results is None:
         st.info(
@@ -294,12 +294,10 @@ def _mini_metric(
     # --------------------------------------------------------
     # Parse expected ToA for comparison
     # --------------------------------------------------------
-    toa_expected: Optional[datetime] = None
+    toa_expected: datetime | None = None
     if stored_config and stored_config.toa_raw:
-        try:
+        with contextlib.suppress(ValueError):
             toa_expected = datetime.strptime(stored_config.toa_raw, "%d/%m/%Y %H:%M:%S")
-        except ValueError:
-            pass
 
     # --------------------------------------------------------
     # Summary KPI cards
@@ -448,7 +446,7 @@ def _mini_metric(
                 label="DBM",
                 color=PLOT_COLORS["DBM"],
             )
-            st.plotly_chart(dbm_fig, use_container_width=True)
+            st.plotly_chart(dbm_fig, width="stretch")
 
     with col_modbm:
         st.markdown(
@@ -587,7 +585,7 @@ def _mini_metric(
                 label="MoDBM",
                 color=PLOT_COLORS["MODBM"],
             )
-            st.plotly_chart(modbm_fig, use_container_width=True)
+            st.plotly_chart(modbm_fig, width="stretch")
 
     # --------------------------------------------------------
     # Full-width combined comparison
@@ -599,4 +597,4 @@ def _mini_metric(
             results=results,
             target_distance_au=stored_config.target.distance,
         )
-        st.plotly_chart(prop_fig, use_container_width=True)
+        st.plotly_chart(prop_fig, width="stretch")

src/heliotrace/models/schemas.py

@@ -11,7 +11,6 @@
 
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional
 
 # ---------------------------------------------------------------------------
 # Input models
@@ -60,9 +59,9 @@ class SimulationConfig:
     c_d: float  # Drag coefficient (dimensionless)
 
     # Optional overrides
-    toa_raw: Optional[str] = None  # Expected arrival "DD/MM/YYYY HH:MM:SS"; None → skip
-    m_override_g: Optional[float] = None  # CME mass override [g]; None → use Pluta formula
-    v0_override_kms: Optional[float] = (
+    toa_raw: str | None = None  # Expected arrival "DD/MM/YYYY HH:MM:SS"; None → skip
+    m_override_g: float | None = None  # CME mass override [g]; None → use Pluta formula
+    v0_override_kms: float | None = (
         None  # Hard override for v₀ toward target [km/s]; None → geometry
     )
 
@@ -109,13 +108,13 @@ class SimulationResults:
     target_hit: bool
 
     # DBM
-    dbm_series: Optional[PropagationSeries]
-    elapsed_time_DBM_h: Optional[float]  # Transit time [h]
-    velocity_arrival_DBM_kms: Optional[float]  # Impact speed [km/s]
-    arrival_time_DBM: Optional[datetime]
+    dbm_series: PropagationSeries | None
+    elapsed_time_DBM_h: float | None  # Transit time [h]
+    velocity_arrival_DBM_kms: float | None  # Impact speed [km/s]
+    arrival_time_DBM: datetime | None
 
     # MoDBM
-    modbm_series: Optional[PropagationSeries]
-    elapsed_time_MODBM_h: Optional[float]
-    velocity_arrival_MODBM_kms: Optional[float]
-    arrival_time_MODBM: Optional[datetime]
+    modbm_series: PropagationSeries | None
+    elapsed_time_MODBM_h: float | None
+    velocity_arrival_MODBM_kms: float | None
+    arrival_time_MODBM: datetime | None

src/heliotrace/physics/drag.py

@@ -13,11 +13,10 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional
 
+import astropy.constants as const
 import astropy.units as u
 import numpy as np
-from astropy.constants import m_p
 from scipy.integrate import OdeSolution, solve_ivp
 from scipy.interpolate import interp1d
 from scipy.optimize import root_scalar
@@ -88,7 +87,7 @@ def get_ambient_solar_wind_density_DBM(r: u.Quantity) -> u.Quantity:
     """
     r_rsun = r.to(u.R_sun).value
     n_init = 3.3e5 * u.cm ** (-3)
-    rho_w = m_p * n_init / r_rsun**2
+    rho_w = const.m_p * n_init / r_rsun**2
     return rho_w
 
 
@@ -97,7 +96,7 @@ def get_constant_drag_parameter_DBM(
     alpha: u.Quantity,
     kappa: float,
     v_apex: u.Quantity,
-    M: Optional[u.Quantity] = None,
+    M: u.Quantity | None = None,
     c_d: float = 1.0,
 ) -> u.Quantity:
     """
@@ -111,10 +110,7 @@ def get_constant_drag_parameter_DBM(
     :param c_d:    Drag coefficient (dimensionless).  Converges to ~1 beyond ~12 R☉.
     :return: Drag parameter γ [1 / km].
     """
-    if M is None:
-        M = get_CME_mass_pluta(v_apex).to(u.kg)
-    else:
-        M = M.to(u.kg)
+    M = get_CME_mass_pluta(v_apex).to(u.kg) if M is None else M.to(u.kg)
 
     rho_w = get_ambient_solar_wind_density_DBM(r)
     A = get_CME_cross_section_pluta(r, alpha, kappa)
@@ -186,7 +182,9 @@ def get_ambient_solar_wind_density_MODBM(r_km: float, ssn: float) -> float:
     :return: Mass density [kg / km³].
     """
     r_au = r_km * u.km.to(u.AU)
-    rho_w = m_p.value * (0.0038 * ssn + 4.5) * (u.cm ** (-3)).to(u.km ** (-3)) * r_au ** (-2.11)
+    rho_w = (
+        const.m_p.value * (0.0038 * ssn + 4.5) * (u.cm ** (-3)).to(u.km ** (-3)) * r_au ** (-2.11)
+    )
     return float(rho_w)
 
 
@@ -265,7 +263,7 @@ def simulate_equation_of_motion_MODBM(
     kappa: float,
     w_type: str,
     ssn: float,
-    M: Optional[u.Quantity] = None,
+    M: u.Quantity | None = None,
     c_d: float = 1.0,
 ) -> OdeSolution:
     """
@@ -285,10 +283,7 @@ def simulate_equation_of_motion_MODBM(
     :return: ``scipy.integrate.solve_ivp`` solution with ``dense_output=True``.
     """
     mass_q: u.Quantity
-    if M is None:
-        mass_q = get_CME_mass_pluta(v_apex).to(u.kg)
-    else:
-        mass_q = M.to(u.kg)
+    mass_q = get_CME_mass_pluta(v_apex).to(u.kg) if M is None else M.to(u.kg)
 
     r0_km = r0.to(u.km).value
     v0_kms = v0.to(u.km / u.s).value

src/heliotrace/physics/fitting.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import logging
-from typing import Optional
 
 import numpy as np
 from numpy.typing import ArrayLike
@@ -17,8 +16,8 @@
 def perform_linear_fit(
     x_data: ArrayLike,
     y_data: ArrayLike,
-    y_error: Optional[float | ArrayLike] = None,
-    initial_guess: Optional[list[float]] = None,
+    y_error: float | ArrayLike | None = None,
+    initial_guess: list[float] | None = None,
 ) -> dict[str, float]:
     """
     Fit a linear model ``y = m * x + b`` to the provided data.

src/heliotrace/ui/components/propagation_plot.py

@@ -6,8 +6,6 @@
 
 from __future__ import annotations
 
-from typing import Optional
-
 import astropy.units as u
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
@@ -29,7 +27,7 @@
 
 def build_single_model_figure(
     series: PropagationSeries,
-    elapsed_h: Optional[float],
+    elapsed_h: float | None,
     target_distance_au: float,
     label: str,
     color: str,
@@ -166,8 +164,8 @@ def build_propagation_comparison_figure(
     )
 
     def _add_model(
-        series: Optional[PropagationSeries],
-        elapsed_h: Optional[float],
+        series: PropagationSeries | None,
+        elapsed_h: float | None,
         label: str,
         color: str,
     ) -> None:

src/heliotrace/ui/components/sidebar_inputs.py

@@ -8,7 +8,6 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Optional
 
 import pandas as pd
 import streamlit as st
@@ -79,7 +78,7 @@ def render_sidebar() -> tuple[SimulationConfig, GCSParams, pd.DataFrame, bool, b
         if st.button(
             "⚡ Fill with example",
             type="primary",
-            use_container_width=True,
+            width="content",
             help="Auto-fill all fields with the 2023-10-28 Halloween CME event.",
         ):
             _fill_example()
@@ -240,7 +239,7 @@ def render_sidebar() -> tuple[SimulationConfig, GCSParams, pd.DataFrame, bool, b
             obs_init,
             hide_index=True,
             num_rows="dynamic",
-            use_container_width=True,
+            width="content",
             column_config={
                 "datetime": st.column_config.DatetimeColumn(
                     "Time (UTC)",
@@ -352,7 +351,7 @@ def render_sidebar() -> tuple[SimulationConfig, GCSParams, pd.DataFrame, bool, b
 
             st.divider()
             st.caption("Optional overrides")
-            toa_raw: Optional[str] = (
+            toa_raw: str | None = (
                 st.text_input(
                     "Expected arrival time (optional)",
                     key="sb_toa_raw",
@@ -377,7 +376,7 @@ def render_sidebar() -> tuple[SimulationConfig, GCSParams, pd.DataFrame, bool, b
                 format="%.2e",
                 help="Override the Pluta (2018) mass formula. Enter 0 to keep the formula.",
             )
-            m_override_g: Optional[float] = m_override_raw if m_override_raw > 0.0 else None
+            m_override_g: float | None = m_override_raw if m_override_raw > 0.0 else None
 
             v0_override_str: str = st.text_input(
                 "v₀ override [km/s]  (blank = geometry-derived)",
@@ -386,7 +385,7 @@ def render_sidebar() -> tuple[SimulationConfig, GCSParams, pd.DataFrame, bool, b
                 help="Hard override for the CME velocity component directed at the target. "
                 r"Replaces the geometry-derived $v_0 = v_{\rm apex} \times$ projection ratio.",
             )
-            v0_override_kms: Optional[float] = _parse_float(v0_override_str)
+            v0_override_kms: float | None = _parse_float(v0_override_str)
             if v0_override_kms is not None:
                 if v0_override_kms <= 0.0:
                     st.error("❌ v₀ override must be > 0 km/s. Override ignored.")
@@ -405,7 +404,7 @@ def render_sidebar() -> tuple[SimulationConfig, GCSParams, pd.DataFrame, bool, b
         run_clicked: bool = st.button(
             "▶ Run Simulation",
             type="primary",
-            use_container_width=True,
+            width="content",
         )
 
     config = SimulationConfig(

src/heliotrace/ui/utils.py

@@ -5,10 +5,9 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Optional
 
 
-def format_diff(sim_time: Optional[datetime], ref_time: Optional[datetime]) -> Optional[str]:
+def format_diff(sim_time: datetime | None, ref_time: datetime | None) -> str | None:
     """
     Format the signed difference between a simulated arrival and a reference time.
 

tests/physics/test_apex_ratio.py

@@ -37,7 +37,7 @@
 
 
 @pytest.mark.parametrize(
-    "cme_lat_deg, cme_lon_deg",
+    ("cme_lat_deg", "cme_lon_deg"),
     [(0, 0), (10, 20), (-5, 15)],
 )
 def test_direct_hit_ratio_near_unity(cme_lat_deg: float, cme_lon_deg: float) -> None:
@@ -60,7 +60,7 @@ def test_direct_hit_ratio_near_unity(cme_lat_deg: float, cme_lon_deg: float) ->
 
 
 @pytest.mark.parametrize(
-    "alpha_deg, kappa",
+    ("alpha_deg", "kappa"),
     [(15, 0.2), (30, 0.42), (45, 0.6)],
 )
 def test_direct_hit_ratio_bounded(alpha_deg: float, kappa: float) -> None:
@@ -232,7 +232,7 @@ def test_tilt_does_not_affect_on_axis_hit(tilt_deg: float) -> None:
 
 
 @pytest.mark.parametrize(
-    "sep_deg, expected_above_threshold",
+    ("sep_deg", "expected_above_threshold"),
     [(30, True), (50, True), (150, False)],
 )
 def test_operationally_relevant_threshold(sep_deg: float, expected_above_threshold: bool) -> None:

tests/physics/test_drag.py

@@ -240,7 +240,7 @@ def test_dbm_gamma_units() -> None:
 
 
 @pytest.mark.parametrize(
-    "r_au, ssn",
+    ("r_au", "ssn"),
     [(1.0, 50.0), (1.5, 100.0), (2.0, 0.0)],
 )
 def test_modbm_density_power_law_exponent(r_au: float, ssn: float) -> None:
@@ -326,7 +326,7 @@ def test_modbm_gamma_inverse_in_mass() -> None:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.parametrize("w_type, expected", [("slow", 363.0), ("fast", 483.0)])
+@pytest.mark.parametrize(("w_type", "expected"), [("slow", 363.0), ("fast", 483.0)])
 def test_modbm_wind_at_1au(w_type: str, expected: float) -> None:
     """At 1 AU, r_AU = 1 → w = w0 * 1.0^0.099 = w0 exactly."""
     w = get_ambient_solar_wind_speed_MODBM(AU_KM, w_type)
@@ -366,7 +366,7 @@ def test_modbm_wind_fast_slow_ratio_constant(r_au: float) -> None:
 @pytest.mark.parametrize("bad_type", ["medium", "", "SLOW", "Fast"])
 def test_modbm_wind_invalid_type_raises(bad_type: str) -> None:
     """Unknown w_type must raise ValueError."""
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="Unknown w_type"):
         get_ambient_solar_wind_speed_MODBM(AU_KM, bad_type)
 
 
@@ -613,7 +613,7 @@ def test_arrival_raises_when_distance_not_reached() -> None:
         400.0 * u.km / u.s,
         1e-7 / u.km,
     )
-    with pytest.raises(ValueError):
+    with pytest.raises(ValueError, match="f\\(a\\) and f\\(b\\) must have different signs"):
         find_time_and_velocity_at_distance(sol, 1.0 * u.AU)
 
 

tests/physics/test_fitting.py

@@ -112,7 +112,7 @@ def test_negative_slope_recovery() -> None:
 
 
 @pytest.mark.parametrize(
-    "N, sigma",
+    ("N", "sigma"),
     [(5, 0.2), (10, 1.0), (20, 0.5), (50, 2.0)],
 )
 def test_weighted_slope_error_closed_form(N: int, sigma: float) -> None: