MSIS Options Analysis Utilities#

Utility functions for creating comprehensive MSIS option analysis plots. This module provides reusable functions for sphinx_gallery documentation.

import matplotlib.pyplot as plt
import numpy as np

import pymsis


def create_option_analysis_figure(option_index, option_name=None, figsize=(14, 10)):
    """
    Create comprehensive analysis plots for a single MSIS option.

    This function generates a 2x2 subplot figure showing:
    1. Altitude profiles with seasonal/diurnal variations
    2. Geographic surface map (latitude vs longitude)
    3. Diurnal cycle (24-hour variation)
    4. Seasonal cycle (annual variation)

    Parameters
    ----------
    option_index : int
        Index of the MSIS option to analyze (0-24)
    option_name : str, optional
        Name of the option for plot titles. If None, uses default names.
    figsize : tuple, optional
        Figure size in inches. Default is (14, 10).

    Returns
    -------
    fig : matplotlib.figure.Figure
        The created figure object
    """

    # Get option name if not provided
    if option_name is None:
        option_names = [
            "F10.7 Effects",
            "Time Independent",
            "Symmetric Annual",
            "Symmetric Semiannual",
            "Asymmetric Annual",
            "Asymmetric Semiannual",
            "Diurnal",
            "Semidiurnal",
            "Geomagnetic Activity",
            "All UT Effects",
            "Longitudinal",
            "Mixed UT/Long",
            "Lower Atmosphere",
            "Lower Atm F10.7",
        ]
        if option_index < len(option_names):
            option_name = option_names[option_index]
        else:
            option_name = f"Option {option_index}"

    # Create figure with 2x2 subplots
    fig, axes = plt.subplots(2, 2, figsize=figsize)

    # Common parameters for all calculations
    f107 = 180  # High solar activity
    f107a = 160  # 81-day average
    ap = 15  # Moderate geomagnetic activity
    aps = [[ap] * 7]  # Ap array format for pymsis

    # Panel 1: Altitude Profiles
    _create_altitude_panel(axes[0, 0], option_index, f107, f107a, aps)

    # Panel 2: Geographic Surface Map
    _create_geographic_panel(axes[0, 1], option_index, f107, f107a, aps)

    # Panel 3: Diurnal Variation
    _create_diurnal_panel(axes[1, 0], option_index, f107, f107a, aps)

    # Panel 4: Seasonal Variation
    _create_seasonal_panel(axes[1, 1], option_index, f107, f107a, aps)

    # Overall title
    fig.suptitle(
        f"Comprehensive Analysis: MSIS {option_name} Option\n"
        f"Atmospheric density changes when option is turned OFF\n"
        f"F10.7={f107}, Ap={ap}",
        fontsize=16,
        y=0.96,
    )

    plt.tight_layout()
    plt.subplots_adjust(top=0.88)

    return fig


def _create_altitude_panel(ax, option_index, f107, f107a, aps):
    """Create altitude profile panel showing seasonal and diurnal effects."""

    # Define four temporal conditions
    date_winter_noon = np.datetime64("2003-01-01T12:00")
    date_winter_midnight = np.datetime64("2003-01-01T00:00")
    date_summer_noon = np.datetime64("2003-07-01T12:00")
    date_summer_midnight = np.datetime64("2003-07-01T00:00")

    lon, lat = 0, 45  # Equator, mid-latitude
    alts = np.linspace(100, 500, 100)  # Thermosphere focus

    # Define options arrays
    options_on = [1] * 25
    options_off = options_on.copy()
    options_off[option_index] = 0

    # Calculate baseline densities (all options ON)
    baselines = {}
    tests = {}

    for condition, date in [
        ("winter_noon", date_winter_noon),
        ("winter_midnight", date_winter_midnight),
        ("summer_noon", date_summer_noon),
        ("summer_midnight", date_summer_midnight),
    ]:
        # Baseline calculation
        baseline = pymsis.calculate(
            date, lon, lat, alts, f107, f107a, aps, options=options_on
        )
        baselines[condition] = np.squeeze(baseline)[:, pymsis.Variable.MASS_DENSITY]

        # Test calculation (option OFF)
        test = pymsis.calculate(
            date, lon, lat, alts, f107, f107a, aps, options=options_off
        )
        tests[condition] = np.squeeze(test)[:, pymsis.Variable.MASS_DENSITY]

    # Calculate and plot percentage differences
    colors = {"winter": "blue", "summer": "red"}
    styles = {"noon": "-", "midnight": "--"}

    for condition, baseline_value in baselines.items():
        season, time = condition.split("_")
        test_value = tests[condition]
        percent_diff = 100 * (test_value - baseline_value) / baseline_value

        label = f"{season.capitalize()} {time}"
        ax.plot(
            percent_diff,
            alts,
            color=colors[season],
            linestyle=styles[time],
            linewidth=2,
            label=label,
        )

    # Formatting
    ax.axvline(x=0, color="gray", linestyle=":", alpha=0.5, linewidth=1)
    ax.set_xlabel("Change when OFF (%)")
    ax.set_ylabel("Altitude (km)")
    ax.set_title("A) Altitude Profiles")
    ax.grid(True, alpha=0.3)
    ax.legend(fontsize=9)


def _create_geographic_panel(ax, option_index, f107, f107a, aps):
    """Create geographic surface map panel (latitude vs longitude)."""

    lats = np.linspace(-90, 90, 19)  # 10-degree spacing
    lons = np.linspace(-180, 180, 37)  # 10-degree spacing
    alt_surface = 300  # km altitude
    date_surface = np.datetime64("2003-06-21T12:00")  # Summer solstice noon

    lon_grid, lat_grid = np.meshgrid(lons, lats)

    # Define options arrays
    options_on = [1] * 25
    options_off = options_on.copy()
    options_off[option_index] = 0

    # Calculate density arrays
    density_on_surface = np.zeros_like(lon_grid)
    density_off_surface = np.zeros_like(lon_grid)

    # Calculate for each lat/lon combination
    for i, lat_val in enumerate(lats):
        for j, lon_val in enumerate(lons):
            # Baseline (ON)
            result_on = pymsis.calculate(
                date_surface,
                lon_val,
                lat_val,
                alt_surface,
                f107,
                f107a,
                aps,
                options=options_on,
            )
            density_on_surface[i, j] = np.squeeze(result_on)[
                pymsis.Variable.MASS_DENSITY
            ]

            # Option OFF
            result_off = pymsis.calculate(
                date_surface,
                lon_val,
                lat_val,
                alt_surface,
                f107,
                f107a,
                aps,
                options=options_off,
            )
            density_off_surface[i, j] = np.squeeze(result_off)[
                pymsis.Variable.MASS_DENSITY
            ]

    # Calculate percentage difference
    percent_diff_surface = (
        100 * (density_off_surface - density_on_surface) / density_on_surface
    )

    # Create contour plot
    levels = np.linspace(-20, 20, 21)
    contour = ax.contourf(
        lon_grid,
        lat_grid,
        percent_diff_surface,
        levels=levels,
        cmap="RdBu_r",
        extend="both",
    )
    ax.contour(
        lon_grid,
        lat_grid,
        percent_diff_surface,
        levels=levels[::4],
        colors="black",
        alpha=0.3,
        linewidths=0.5,
    )

    # Formatting
    ax.set_xlabel("Longitude (°)")
    ax.set_ylabel("Latitude (°)")
    ax.set_title(f"B) Geographic Pattern at {alt_surface} km")
    ax.grid(True, alpha=0.3)

    # Add colorbar
    cbar = plt.colorbar(contour, ax=ax, shrink=0.8)
    cbar.set_label("Change when OFF (%)", fontsize=9)


def _create_diurnal_panel(ax, option_index, f107, f107a, aps):
    """Create diurnal cycle panel (24-hour variation)."""

    times = np.linspace(0, 24, 25)
    dates_temporal = [
        np.datetime64("2003-06-21") + np.timedelta64(int(h), "h") for h in times
    ]

    alt_temporal = 300
    lat_temporal = 45
    lon_temporal = 0

    # Define options arrays
    options_on = [1] * 25
    options_off = options_on.copy()
    options_off[option_index] = 0

    density_on_temporal = []
    density_off_temporal = []

    for date in dates_temporal:
        # Baseline (ON)
        result_on = pymsis.calculate(
            date,
            lon_temporal,
            lat_temporal,
            alt_temporal,
            f107,
            f107a,
            aps,
            options=options_on,
        )
        density_on_temporal.append(np.squeeze(result_on)[pymsis.Variable.MASS_DENSITY])

        # Option OFF
        result_off = pymsis.calculate(
            date,
            lon_temporal,
            lat_temporal,
            alt_temporal,
            f107,
            f107a,
            aps,
            options=options_off,
        )
        density_off_temporal.append(
            np.squeeze(result_off)[pymsis.Variable.MASS_DENSITY]
        )

    density_on_temporal = np.array(density_on_temporal)
    density_off_temporal = np.array(density_off_temporal)

    percent_diff_temporal = (
        100 * (density_off_temporal - density_on_temporal) / density_on_temporal
    )

    # Plot
    ax.plot(times, percent_diff_temporal, "b-", linewidth=2, marker="o", markersize=4)
    ax.axhline(y=0, color="gray", linestyle=":", alpha=0.5, linewidth=1)

    # Formatting
    ax.set_xlabel("Time (hours UT)")
    ax.set_ylabel("Change when OFF (%)")
    ax.set_title(f"C) Diurnal Cycle at {lat_temporal}°N, {alt_temporal} km")
    ax.grid(True, alpha=0.3)
    ax.set_xlim(0, 24)


def _create_seasonal_panel(ax, option_index, f107, f107a, aps):
    """Create seasonal variation panel (annual cycle)."""

    # Create dates for full year (monthly sampling)
    months = np.arange(1, 13)
    dates_seasonal = [np.datetime64(f"2003-{month:02d}-15T12:00") for month in months]

    alt_temporal = 300
    lat_temporal = 45
    lon_temporal = 0

    # Define options arrays
    options_on = [1] * 25
    options_off = options_on.copy()
    options_off[option_index] = 0

    density_on_seasonal = []
    density_off_seasonal = []

    for date in dates_seasonal:
        # Baseline (ON)
        result_on = pymsis.calculate(
            date,
            lon_temporal,
            lat_temporal,
            alt_temporal,
            f107,
            f107a,
            aps,
            options=options_on,
        )
        density_on_seasonal.append(np.squeeze(result_on)[pymsis.Variable.MASS_DENSITY])

        # Option OFF
        result_off = pymsis.calculate(
            date,
            lon_temporal,
            lat_temporal,
            alt_temporal,
            f107,
            f107a,
            aps,
            options=options_off,
        )
        density_off_seasonal.append(
            np.squeeze(result_off)[pymsis.Variable.MASS_DENSITY]
        )

    density_on_seasonal = np.array(density_on_seasonal)
    density_off_seasonal = np.array(density_off_seasonal)

    percent_diff_seasonal = (
        100 * (density_off_seasonal - density_on_seasonal) / density_on_seasonal
    )

    # Plot
    ax.plot(months, percent_diff_seasonal, "r-", linewidth=2, marker="s", markersize=6)
    ax.axhline(y=0, color="gray", linestyle=":", alpha=0.5, linewidth=1)

    # Formatting
    ax.set_xlabel("Month")
    ax.set_ylabel("Change when OFF (%)")
    ax.set_title(f"D) Seasonal Cycle at {lat_temporal}°N, {alt_temporal} km")
    ax.grid(True, alpha=0.3)
    ax.set_xlim(1, 12)
    ax.set_xticks(months)

Gallery generated by Sphinx-Gallery