pycontrails 0.51.1__cp311-cp311-macosx_11_0_arm64.whl → 0.52.0__cp311-cp311-macosx_11_0_arm64.whl

pycontrails/models/apcemm/utils.py

@@ -0,0 +1,437 @@
+"""APCEMM interface utility functions."""
+
+from __future__ import annotations
+
+import pathlib
+import subprocess
+from typing import Any
+
+import numpy as np
+import xarray as xr
+
+from pycontrails.core import GeoVectorDataset, MetDataset, met_var, models
+from pycontrails.models.apcemm.inputs import APCEMMInput
+from pycontrails.models.humidity_scaling import HumidityScaling
+from pycontrails.physics import constants, thermo, units
+from pycontrails.utils.types import ArrayScalarLike
+
+_path_to_static = pathlib.Path(__file__).parent / "static"
+YAML_TEMPLATE = _path_to_static / "apcemm_yaml_template.yaml"
+
+
+def generate_apcemm_input_yaml(params: APCEMMInput) -> str:
+    """Generate YAML file from  APCEMM input parameters.
+
+    Parameters
+    ----------
+    params : APCEMMInput
+        :class:`APCEMMInput` instance with parameters for input YAML file.
+
+    Return
+    ------
+    str
+        Contents of input YAML file generated from :param:`params`.
+    """
+
+    with open(YAML_TEMPLATE) as f:
+        template = f.read()
+
+    return template.format(
+        n_threads_int=params.n_threads,
+        output_folder_str=params.output_directory,
+        overwrite_output_bool=_yaml_bool(params.overwrite_output),
+        input_background_conditions_str=params.input_background_conditions,
+        input_engine_emissions_str=params.input_engine_emissions,
+        max_age_hr=params.max_age / np.timedelta64(1, "h"),
+        temperature_k=params.air_temperature,
+        rhw_percent=params.rhw * 100,
+        horiz_diff_coeff_m2_per_s=params.horiz_diff,
+        vert_diff_coeff_m2_per_s=params.vert_diff,
+        pressure_hpa=params.air_pressure / 100,
+        wind_shear_per_s=params.normal_shear,
+        brunt_vaisala_frequency_per_s=params.brunt_vaisala_frequency,
+        longitude_deg=params.longitude,
+        latitude_deg=params.latitude,
+        emission_day_dayofyear=params.day_of_year,
+        emission_time_hourofday=params.hour_of_day,
+        nox_ppt=params.nox_vmr * 1e12,
+        hno3_ppt=params.hno3_vmr * 1e12,
+        o3_ppb=params.o3_vmr * 1e9,
+        co_ppb=params.co_vmr * 1e9,
+        ch4_ppm=params.ch4_vmr * 1e6,
+        so2_ppt=params.so2_vmr * 1e12,
+        nox_g_per_kg=params.nox_ei * 1000,
+        co_g_per_kg=params.co_ei * 1000,
+        uhc_g_per_kg=params.hc_ei * 1000,
+        so2_g_per_kg=params.so2_ei * 1000,
+        so2_to_so4_conv_percent=params.so2_to_so4_conversion * 100,
+        soot_g_per_kg=params.nvpm_ei_m * 1000,
+        soot_radius_m=params.soot_radius,
+        total_fuel_flow_kg_per_s=params.fuel_flow,
+        aircraft_mass_kg=params.aircraft_mass,
+        flight_speed_m_per_s=params.true_airspeed,
+        num_of_engines_int=params.n_engine,
+        wingspan_m=params.wingspan,
+        core_exit_temp_k=params.core_exit_temp,
+        exit_bypass_area_m2=params.core_exit_area,
+        transport_timestep_min=params.dt_apcemm_transport / np.timedelta64(1, "m"),
+        gravitational_settling_bool=_yaml_bool(params.do_gravitational_setting),
+        solid_coagulation_bool=_yaml_bool(params.do_solid_coagulation),
+        liquid_coagulation_bool=_yaml_bool(params.do_liquid_coagulation),
+        ice_growth_bool=_yaml_bool(params.do_ice_growth),
+        coag_timestep_min=params.dt_apcemm_coagulation / np.timedelta64(1, "m"),
+        ice_growth_timestep_min=params.dt_apcemm_ice_growth / np.timedelta64(1, "m"),
+        met_input_file_path_str=params.input_met_file,
+        time_series_data_timestep_hr=params.dt_input_met / np.timedelta64(1, "h"),
+        save_aerosol_timeseries_bool=_yaml_bool(params.do_apcemm_nc_output),
+        aerosol_indices_list_int=",".join(str(i) for i in params.apcemm_nc_output_species),
+        save_frequency_min=params.dt_apcemm_nc_output / np.timedelta64(1, "m"),
+        nx_pos_int=params.nx,
+        ny_pos_int=params.ny,
+        xlim_right_pos_m=params.xlim_right,
+        xlim_left_pos_m=params.xlim_left,
+        ylim_up_pos_m=params.ylim_up,
+        ylim_down_pos_m=params.ylim_down,
+        base_contrail_depth_m=params.initial_contrail_depth_offset,
+        contrail_depth_scaling_factor_nondim=params.initial_contrail_depth_scale_factor,
+        base_contrail_width_m=params.initial_contrail_width_offset,
+        contrail_width_scaling_factor_nondim=params.initial_contrail_width_scale_factor,
+    )
+
+
+def _yaml_bool(param: bool) -> str:
+    """Convert boolean to T/F for YAML file."""
+    return "T" if param else "F"
+
+
+def generate_apcemm_input_met(
+    time: np.ndarray,
+    longitude: np.ndarray,
+    latitude: np.ndarray,
+    azimuth: np.ndarray,
+    altitude: np.ndarray,
+    met: MetDataset,
+    humidity_scaling: HumidityScaling,
+    dz_m: float,
+    interp_kwargs: dict[str, Any],
+) -> xr.Dataset:
+    r"""Create xarray Dataset for APCEMM meteorology netCDF file.
+
+    This dataset contains a sequence of atmospheric profiles along the
+    Lagrangian trajectory of an advected flight segment. The along-trajectory
+    dimension is parameterized by time (rather than latitude and longitude),
+    so the dataset coordinates are air pressure and time.
+
+    Parameters
+    ----------
+    time : np.ndarray
+        Time coordinates along the Lagrangian trajectory of the advected flight segment.
+        Values must be coercible to ``np.datetime64`` by :class:`GeoVectorDataset`.
+        Will be flattened before use if not 1-dimensional.
+    longitude : np.ndarray
+        Longitude [WGS84] along the Lagrangian trajectory of the advected flight segment.
+        Defines the longitude of the trajectory at each time and should have the
+        same shape as :param:`time`
+        Will be flattened before use if not 1-dimensional.
+    latitude : np.ndarray
+        Latitude [WGS84] along the Lagrangian trajectory of the advected flight segment.
+        Defines the longitude of the trajectory at each time and should have the
+        same shape as :param:`time`
+        Will be flattened before use if not 1-dimensional.
+    azimuth : np.ndarray
+        Azimuth [:math:`\deg`] of the advected flight segment at each point along its
+        Lagrangian trajectory. Note that the azimuth defines the orientation of the
+        advected segment itself, and not the direction in which advection is transporting
+        the segment. The azimuth is used to convert horizontal winds into segment-normal
+        wind shear. Must have the same shape as :param:`time`.
+        Will be flattened before use if not 1-dimensional.
+    altitude : np.ndarray
+        Defines altitudes [:math:`m`] on which atmospheric profiles are computed.
+        Profiles are defined using the same set of altitudes at every point
+        along the Lagrangian trajectory of the advected flight segment. Note that
+        this parameter does not have to have the same shape as :param:`time`.
+    met : MetDataset
+        Meteorology used to generate the sequence of atmospheric profiles. Must contain:
+        - air temperature [:math:`K`]
+        - specific humidity [:math:`kg/kg`]
+        - geopotential height [:math:`m`]
+        - eastward wind [:math:`m/s`]
+        - northward wind [:math:`m/s`]
+        - vertical velocity [:math:`Pa/s`]
+    humidity_scaling : HumidityScaling
+        Humidity scaling applied to specific humidity in :param:`met` before
+        generating atmospheric profiles.
+    dz_m : float
+        Altitude difference [:math:`m`] used to approximate vertical derivatives
+        when computing wind shear.
+
+    Returns
+    -------
+    xr.Dataset
+        Meteorology dataset in required format for APCEMM input.
+    """
+
+    # Ensure that altitudes are sorted ascending
+    altitude = np.sort(altitude)
+
+    # Check for required fields in met
+    vars = (
+        met_var.AirTemperature,
+        met_var.SpecificHumidity,
+        met_var.GeopotentialHeight,
+        met_var.EastwardWind,
+        met_var.NorthwardWind,
+        met_var.VerticalVelocity,
+    )
+    met.ensure_vars(vars)
+    met.standardize_variables(vars)
+
+    # Flatten input arrays
+    time = time.ravel()
+    longitude = longitude.ravel()
+    latitude = latitude.ravel()
+    azimuth = azimuth.ravel()
+    altitude = altitude.ravel()
+
+    # Estimate pressure levels close to target altitudes
+    # (not exact because this assumes the ISA temperature profile)
+    pressure = units.m_to_pl(altitude) * 1e2
+
+    # Broadcast to required shape and create vector for initial interpolation
+    # onto original pressure levels at target horizontal location.
+    shape = (time.size, altitude.size)
+    time = np.broadcast_to(time[:, np.newaxis], shape).ravel()
+    longitude = np.broadcast_to(longitude[:, np.newaxis], shape).ravel()
+    latitude = np.broadcast_to(latitude[:, np.newaxis], shape).ravel()
+    azimuth = np.broadcast_to(azimuth[:, np.newaxis], shape).ravel()
+    level = np.broadcast_to(pressure[np.newaxis, :] / 1e2, shape).ravel()
+    vector = GeoVectorDataset(
+        data={"azimuth": azimuth},
+        longitude=longitude,
+        latitude=latitude,
+        level=level,
+        time=time,
+    )
+
+    # Downselect met before interpolation
+    met = vector.downselect_met(met, copy=False)
+
+    # Interpolate meteorology data onto vector
+    scale_humidity = humidity_scaling is not None and "specific_humidity" not in vector
+    for met_key in (
+        "air_temperature",
+        "eastward_wind",
+        "geopotential_height",
+        "northward_wind",
+        "specific_humidity",
+        "lagrangian_tendency_of_air_pressure",
+    ):
+        models.interpolate_met(met, vector, met_key, **interp_kwargs)
+
+    # Interpolate winds at lower level for shear calculation
+    air_pressure_lower = thermo.pressure_dz(vector["air_temperature"], vector.air_pressure, dz_m)
+    lower_level = air_pressure_lower / 100.0
+    for met_key in ("eastward_wind", "northward_wind"):
+        vector_key = f"{met_key}_lower"
+        models.interpolate_met(met, vector, met_key, vector_key, **interp_kwargs, level=lower_level)
+
+    # Apply humidity scaling
+    if scale_humidity and humidity_scaling is not None:
+        humidity_scaling.eval(vector, copy_source=False)
+
+    # Compute RHi and segment-normal shear
+    vector.setdefault(
+        "rhi",
+        thermo.rhi(vector["specific_humidity"], vector["air_temperature"], vector.air_pressure),
+    )
+    vector.setdefault(
+        "normal_shear",
+        normal_wind_shear(
+            vector["eastward_wind"],
+            vector["eastward_wind_lower"],
+            vector["northward_wind"],
+            vector["northward_wind_lower"],
+            vector["azimuth"],
+            dz_m,
+        ),
+    )
+
+    # Reshape interpolated fields to (time, level).
+    nlev = altitude.size
+    ntime = len(vector) // nlev
+    shape = (ntime, nlev)
+    time = np.unique(vector["time"])
+    time = (time - time[0]) / np.timedelta64(1, "h")
+    temperature = vector["air_temperature"].reshape(shape)
+    qv = vector["specific_humidity"].reshape(shape)
+    z = vector["geopotential_height"].reshape(shape)
+    rhi = vector["rhi"].reshape(shape)
+    shear = vector["normal_shear"].reshape(shape)
+    shear[:, -1] = shear[:, -2]  # lowest level will be nan
+    omega = vector["lagrangian_tendency_of_air_pressure"].reshape(shape)
+    virtual_temperature = temperature * (1 + qv / constants.epsilon) / (1 + qv)
+    density = pressure[np.newaxis, :] / (constants.R_d * virtual_temperature)
+    w = -omega / (density * constants.g)
+
+    # Interpolate fields to target altitudes profile-by-profile
+    # to obtain 2D arrays with dimensions (time, altitude).
+    temperature_on_z = np.zeros(shape, dtype=temperature.dtype)
+    rhi_on_z = np.zeros(shape, dtype=rhi.dtype)
+    shear_on_z = np.zeros(shape, dtype=shear.dtype)
+    w_on_z = np.zeros(shape, dtype=w.dtype)
+
+    # Fields should already be on pressure levels close to target
+    # altitudes, so this just uses linear interpolation and constant
+    # extrapolation on fields expected by APCEMM.
+    # NaNs are preserved at the start and end of interpolated profiles
+    # but removed in interiors.
+    def interp(z: np.ndarray, z0: np.ndarray, f0: np.ndarray) -> np.ndarray:
+        # mask nans
+        mask = np.isnan(z0) | np.isnan(f0)
+        if np.all(mask):
+            msg = (
+                "Found all-NaN profile during APCEMM meterology input file creation. "
+                "MetDataset may have insufficient spatiotemporal coverage."
+            )
+            raise ValueError(msg)
+        z0 = z0[~mask]
+        f0 = f0[~mask]
+
+        # interpolate
+        assert np.all(np.diff(z0) > 0)  # expect increasing altitudes
+        fi = np.interp(z, z0, f0, left=f0[0], right=f0[-1])
+
+        # restore nans at start and end of profile
+        if mask[0]:  # nans at top of profile
+            fi[z > z0.max()] = np.nan
+        if mask[-1]:  # nans at end of profile
+            fi[z < z0.min()] = np.nan
+        return fi
+
+    # The manual for loop is unlikely to be a bottleneck since a
+    # substantial amount of work is done within each iteration.
+    for i in range(ntime):
+        temperature_on_z[i, :] = interp(altitude, z[i, :], temperature[i, :])
+        rhi_on_z[i, :] = interp(altitude, z[i, :], rhi[i, :])
+        shear_on_z[i, :] = interp(altitude, z[i, :], shear[i, :])
+        w_on_z[i, :] = interp(altitude, z[i, :], w[i, :])
+
+    # APCEMM also requires initial pressure profile
+    pressure_on_z = interp(altitude, z[0, :], pressure)
+
+    # Create APCEMM input dataset.
+    # Transpose require because APCEMM expects (altitude, time) arrays.
+    return xr.Dataset(
+        data_vars={
+            "pressure": (("altitude",), pressure_on_z.astype("float32") / 1e2, {"units": "hPa"}),
+            "temperature": (
+                ("altitude", "time"),
+                temperature_on_z.astype("float32").T,
+                {"units": "K"},
+            ),
+            "relative_humidity_ice": (
+                ("altitude", "time"),
+                1e2 * rhi_on_z.astype("float32").T,
+                {"units": "percent"},
+            ),
+            "shear": (("altitude", "time"), shear_on_z.astype("float32").T, {"units": "s**-1"}),
+            "w": (("altitude", "time"), w_on_z.astype("float32").T, {"units": "m s**-1"}),
+        },
+        coords={
+            "altitude": ("altitude", altitude.astype("float32") / 1e3, {"units": "km"}),
+            "time": ("time", time, {"units": "hours"}),
+        },
+    )
+
+
+def run(
+    apcemm_path: pathlib.Path | str, input_yaml: str, rundir: str, stdout_log: str, stderr_log: str
+) -> None:
+    """
+    Run APCEMM executable.
+
+    Parameters
+    ----------
+    apcemm_path : pathlib.Path | str
+        Path to APCEMM executable.
+    input_yaml : str
+        Path to APCEMM input yaml file.
+    rundir : str
+        Path to APCEMM simulation directory.
+    stdout_log : str
+        Path to file used to log APCEMM stdout
+    stderr_log  : str
+        Path to file used to log APCEMM stderr
+
+    Raises
+    ------
+    ChildProcessError
+        APCEMM exits with a non-zero return code.
+    """
+
+    with open(stdout_log, "w") as stdout, open(stderr_log, "w") as stderr:
+        result = subprocess.run(
+            [apcemm_path, input_yaml], stdout=stdout, stderr=stderr, cwd=rundir, check=False
+        )
+        if result.returncode != 0:
+            msg = (
+                f"APCEMM simulation in {rundir} "
+                f"exited with return code {result.returncode}. "
+                f"Check logs at {stdout_log} and {stderr_log}."
+            )
+            raise ChildProcessError(msg)
+
+
+def normal_wind_shear(
+    u_hi: ArrayScalarLike,
+    u_lo: ArrayScalarLike,
+    v_hi: ArrayScalarLike,
+    v_lo: ArrayScalarLike,
+    azimuth: ArrayScalarLike,
+    dz: float,
+) -> ArrayScalarLike:
+    r"""Compute segment-normal wind shear from wind speeds at lower and upper levels.
+
+    Parameters
+    ----------
+    u_hi : ArrayScalarLike
+        Eastward wind at upper level [:math:`m/s`]
+    u_lo : ArrayScalarLike
+        Eastward wind at lower level [:math:`m/s`]
+    v_hi : ArrayScalarLike
+        Northward wind at upper level [:math:`m/s`]
+    v_lo : ArrayScalarLike
+        Northward wind at lower level [:math:`m/s`]
+    azimuth : ArrayScalarLike
+        Segment azimuth [:math:`\deg`]
+    dz : float
+        Distance between upper and lower level [:math:`m`]
+
+    Returns
+    -------
+    ArrayScalarLike
+        Segment-normal wind shear [:math:`1/s`]
+    """
+    du_dz = (u_hi - u_lo) / dz
+    dv_dz = (v_hi - v_lo) / dz
+    az_radians = units.degrees_to_radians(azimuth)
+    sin_az = np.sin(az_radians)
+    cos_az = np.cos(az_radians)
+    return sin_az * dv_dz - cos_az * du_dz
+
+
+def soot_radius(
+    nvpm_ei_m: ArrayScalarLike, nvpm_ei_n: ArrayScalarLike, rho_bc: float = 1770.0
+) -> ArrayScalarLike:
+    """Calculate mean soot radius from mass and number emissions indices.
+
+    Parameters
+    ----------
+    nvpm_ei_m : ArrayScalarLike
+        Soot mass emissions index [:math:`kg/kg`]
+    nvpm_ei_n : ArrayScalarLike
+        Soot number emissions index [:math:`1/kg`]
+    rho_bc : float, optional
+        Density of black carbon [:math:`kg/m^3`]. By default, 1770.
+    """
+    return ((3.0 * nvpm_ei_m) / (4.0 * np.pi * rho_bc * nvpm_ei_n)) ** (1.0 / 3.0)

pycontrails/models/cocip/__init__.py

@@ -4,6 +4,7 @@ from pycontrails.models.cocip.cocip import Cocip
 from pycontrails.models.cocip.cocip_params import CocipFlightParams, CocipParams
 from pycontrails.models.cocip.cocip_uncertainty import CocipUncertaintyParams, habit_dirichlet
 from pycontrails.models.cocip.output_formats import (
+    compare_cocip_with_goes,
     contrail_flight_summary_statistics,
     contrails_to_hi_res_grid,
     flight_waypoint_summary_statistics,
@@ -24,4 +25,5 @@ __all__ = [
     "longitude_latitude_grid",
     "natural_cirrus_properties_to_hi_res_grid",
     "time_slice_statistics",
+    "compare_cocip_with_goes",
 ]

pycontrails/models/cocip/output_formats.py

@@ -14,13 +14,17 @@ This module includes functions to produce additional output formats, including t
     (6) Increase spatial resolution of natural cirrus properties, required to estimate the
         high-resolution contrail cirrus coverage for (5).
             See :func:`natural_cirrus_properties_to_hi_res_grid`.
+    (7) Comparing simulated contrails from CoCiP with GOES satellite imagery.
+            See :func:`compare_cocip_with_goes`.
 """
 
 from __future__ import annotations
 
+import pathlib
 import warnings
 from collections.abc import Hashable
 
+import matplotlib.pyplot as plt
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
@@ -28,6 +32,7 @@ import xarray as xr
 
 from pycontrails.core.met import MetDataArray, MetDataset
 from pycontrails.core.vector import GeoVectorDataset, vector_to_lon_lat_grid
+from pycontrails.datalib.goes import GOES, extract_goes_visualization
 from pycontrails.models.cocip.contrail_properties import contrail_edges, plume_mass_per_distance
 from pycontrails.models.cocip.radiative_forcing import albedo
 from pycontrails.models.humidity_scaling import HumidityScaling
@@ -2077,3 +2082,163 @@ def _repeat_rows_and_columns(
 
     # Do not repeat final row and column as they are on the edge
     return array_2d_rep[: -(n_reps - 1), : -(n_reps - 1)]
+
+
+# -----------------------------------------
+# Compare CoCiP outputs with GOES satellite
+# -----------------------------------------
+
+
+def compare_cocip_with_goes(
+    time: np.timedelta64 | pd.Timestamp,
+    flight: GeoVectorDataset | pd.DataFrame,
+    contrail: GeoVectorDataset | pd.DataFrame,
+    *,
+    spatial_bbox: tuple[float, float, float, float] = (-160.0, -80.0, 10.0, 80.0),
+    region: str = "F",
+    path_write_img: pathlib.Path | None = None,
+) -> None | pathlib.Path:
+    r"""
+    Compare simulated persistent contrails from CoCiP with GOES satellite imagery.
+
+    Parameters
+    ----------
+    time : np.timedelta64 | pd.Timestamp
+        Time of GOES satellite image.
+    flight : GeoVectorDataset | pd.DataFrame
+        Flight waypoints.
+        Best to use the returned output :class:`Flight` from
+        :meth:`pycontrails.models.cocip.Cocip.eval`.
+    contrail : GeoVectorDataset | pd.DataFrame,
+        Contrail evolution outputs (:attr:`pycontrails.models.cocip.Cocip.contrail`)
+        set during :meth:`pycontrails.models.cocip.Cocip.eval`.
+    spatial_bbox : tuple[float, float, float, float]
+        Spatial bounding box, ``(lon_min, lat_min, lon_max, lat_max)``, [:math:`\deg`]
+    region : str
+        'F' for full disk (image provided every 10 m), and 'C' for CONUS (image provided every 5 m)
+    path_write_img : None | pathlib.Path
+        File path to save the CoCiP-GOES image.
+
+    Returns
+    -------
+    None | pathlib.Path
+        File path of saved CoCiP-GOES image if ``path_write_img`` is provided.
+    """
+
+    try:
+        import cartopy.crs as ccrs
+        from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter
+    except ModuleNotFoundError as e:
+        dependencies.raise_module_not_found_error(
+            name="compare_cocip_with_goes function",
+            package_name="cartopy",
+            module_not_found_error=e,
+            pycontrails_optional_package="goes",
+        )
+
+    # Round `time` to nearest GOES image time slice
+    if isinstance(time, np.timedelta64):
+        time = pd.to_datetime(time)
+
+    if region == "F":
+        time = time.round("10min")
+    elif region == "C":
+        time = time.round("5min")
+    else:
+        raise AssertionError("`region` only accepts inputs of `F` (full disk) or `C` (CONUS)")
+
+    _flight = GeoVectorDataset(flight)
+    _contrail = GeoVectorDataset(contrail)
+
+    # Ensure the required columns are included in `flight_waypoints` and `contrails`
+    _flight.ensure_vars(["flight_id", "waypoint"])
+    _contrail.ensure_vars(
+        ["flight_id", "waypoint", "sin_a", "cos_a", "width", "tau_contrail", "age_hours"]
+    )
+
+    # Downselect `_flight` only to spatial domain covered by GOES full disk
+    is_in_lon = _flight.dataframe["longitude"].between(spatial_bbox[0], spatial_bbox[2])
+    is_in_lat = _flight.dataframe["latitude"].between(spatial_bbox[1], spatial_bbox[3])
+    is_in_lon_lat = is_in_lon & is_in_lat
+
+    if not np.any(is_in_lon_lat):
+        warnings.warn(
+            "Flight trajectory does not intersect with the defined spatial bounding box or spatial "
+            "domain covered by GOES."
+        )
+
+    _flight = _flight.filter(is_in_lon_lat)
+
+    # Filter `_flight` if time bounds were previously defined.
+    is_before_time = _flight["time"] < time
+
+    if not np.any(is_before_time):
+        warnings.warn("No flight waypoints were recorded before the specified `time`.")
+
+    _flight = _flight.filter(is_before_time)
+
+    # Downselect `_contrail` only to include the filtered flight waypoints
+    is_in_domain = _contrail.dataframe["waypoint"].isin(_flight["waypoint"])
+
+    if not np.any(is_in_domain):
+        warnings.warn(
+            "No persistent contrails were formed within the defined spatial bounding box."
+        )
+
+    _contrail = _contrail.filter(is_in_domain)
+
+    # Download GOES image at `time`
+    goes = GOES(region=region)
+    da = goes.get(time)
+    rgb, transform, extent = extract_goes_visualization(da)
+    bbox = spatial_bbox[0], spatial_bbox[2], spatial_bbox[1], spatial_bbox[3]
+
+    # Calculate optimal figure dimensions
+    d_lon = spatial_bbox[2] - spatial_bbox[0]
+    d_lat = spatial_bbox[3] - spatial_bbox[1]
+    x_dim = 9.99
+    y_dim = x_dim * (d_lat / d_lon)
+
+    # Plot data
+    fig = plt.figure(figsize=(1.2 * x_dim, y_dim))
+    pc = ccrs.PlateCarree()
+    ax = fig.add_subplot(projection=pc, extent=bbox)
+    ax.coastlines()  # type: ignore[attr-defined]
+    ax.imshow(rgb, extent=extent, transform=transform)
+
+    ax.set_xticks([spatial_bbox[0], spatial_bbox[2]], crs=ccrs.PlateCarree())
+    ax.set_yticks([spatial_bbox[1], spatial_bbox[3]], crs=ccrs.PlateCarree())
+    lon_formatter = LongitudeFormatter(zero_direction_label=True)
+    lat_formatter = LatitudeFormatter()
+    ax.xaxis.set_major_formatter(lon_formatter)
+    ax.yaxis.set_major_formatter(lat_formatter)
+
+    # Plot flight trajectory up to `time`
+    ax.plot(_flight["longitude"], _flight["latitude"], c="k", linewidth=2.5)
+    plt.legend(["Flight trajectory"])
+
+    # Plot persistent contrails at `time`
+    is_time = (_contrail["time"] == time) & (~np.isnan(_contrail["age_hours"]))
+    im = ax.scatter(
+        _contrail["longitude"][is_time],
+        _contrail["latitude"][is_time],
+        c=_contrail["tau_contrail"][is_time],
+        s=4,
+        cmap="YlOrRd_r",
+        vmin=0,
+        vmax=0.2,
+    )
+    cbar = plt.colorbar(im)
+    cbar.set_label(r"$\tau_{\rm contrail}$")
+    ax.set_title(f"{time}")
+    plt.tight_layout()
+
+    # return output path if `path_write_img` is not None
+    if path_write_img is not None:
+        t_str = time.strftime("%Y%m%d_%H%M%S")
+        file_name = f"goes_{t_str}.png"
+        output_path = path_write_img.joinpath(file_name)
+        plt.savefig(output_path, dpi=150, bbox_inches="tight")
+        plt.close()
+
+        return output_path

pycontrails/models/cocipgrid/cocip_grid.py

@@ -749,12 +749,13 @@ class CocipGrid(models.Model):
 
         # These should probably not be included in model input ... so
         # we'll get a warning if they get overwritten
-        vector["longitude_head"], vector["latitude_head"] = geo.forward_azimuth(
-            lons=lons, lats=lats, az=azimuth, dist=dist
-        )
-        vector["longitude_tail"], vector["latitude_tail"] = geo.forward_azimuth(
-            lons=lons, lats=lats, az=azimuth, dist=-dist
-        )
+        lon_head, lat_head = geo.forward_azimuth(lons=lons, lats=lats, az=azimuth, dist=dist)
+        vector["longitude_head"] = lon_head.astype(self._target_dtype, copy=False)
+        vector["latitude_head"] = lat_head.astype(self._target_dtype, copy=False)
+
+        lon_tail, lat_tail = geo.forward_azimuth(lons=lons, lats=lats, az=azimuth, dist=-dist)
+        vector["longitude_tail"] = lon_tail.astype(self._target_dtype, copy=False)
+        vector["latitude_tail"] = lat_tail.astype(self._target_dtype, copy=False)
 
         return vector