import logging
from typing import List, Literal, Optional, Tuple
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
import quivr as qv
from ..constants import KM_P_AU
from ..constants import Constants as c
from ..coordinates import (
CartesianCoordinates,
GeodeticCoordinates,
Origin,
OriginCodes,
SphericalCoordinates,
transform_coordinates,
)
from ..orbits import Orbits, VariantOrbits
from ..propagator import Propagator
from ..time import Timestamp
from ..utils.plots.data import Coastlines
from ..utils.plots.logos import (
AsteroidInstituteLogoDark,
AsteroidInstituteLogoLight,
get_logo_base64,
)
from ..utils.spice import get_perturber_state
from .impacts import CollisionEvent
EARTH_RADIUS_KM = c.R_EARTH_EQUATORIAL * KM_P_AU * 0.999
MOON_RADIUS_KM = 1738.1
logger = logging.getLogger(__name__)
try:
import geopandas as gpd
import plotly.graph_objects as go
except ImportError:
raise ImportError("Please install adam_core[plots] to use this feature.")
[docs]
def prepare_propagated_variants(
propagated_variants: Orbits, impacts: CollisionEvent
) -> dict[Literal["Non-Impacting", "EARTH", "MOON"], Orbits]:
"""
Sets variants propagated after their impact time to their impact coordinates on the surface of the colliding body. If the colliding body is the Earth,
the variants are set to the impact coordinates on the surface of the Earth. If the colliding body is the Moon, the variants are set to the impact coordinates
in the lunarcentric frame (but not fixed to the surface of the Moon).
Note: Due to the nature of the impact detection code, some of the variants may already be inside the sphere of the Earth when the impact is detected.
In these cases, the variants' distance from the geocenter is set to the radius of the Earth. The results of this function should not be used
for high-fidelity impact predictions but instead for visualizations of the approximate impact corridor.
Parameters
----------
propagated_variants: Orbits
The propagated variants to cleanse.
impacts: CollisionEvent
The impacts detected within the variants.
Returns
-------
dict[Literal["Non-Impacting", "EARTH", "MOON"], Orbits]
A dictionary containing the prepared variants, with keys:
- "Non-Impacting": Variants that don't impact any body
- "EARTH": Variants that impact Earth (if any)
- "MOON": Variants that impact the Moon (if any)
Only bodies that appear in the impacts will be included as keys.
"""
assert propagated_variants.coordinates.frame == "ecliptic"
colliding_bodies = impacts.collision_object.code.unique().to_pylist()
prepared_variants = {}
# Remove the non-impacting variants
impacting_variants = propagated_variants.apply_mask(
pc.is_in(propagated_variants.orbit_id, impacts.variant_id)
)
non_impacting_variants = propagated_variants.apply_mask(
pc.invert(pc.is_in(propagated_variants.orbit_id, impacts.variant_id))
)
prepared_variants["Non-Impacting"] = non_impacting_variants
for colliding_body in colliding_bodies:
if colliding_body == "EARTH":
radius = EARTH_RADIUS_KM
elif colliding_body == "MOON":
radius = MOON_RADIUS_KM
else:
raise ValueError(
f"CollisionEvent visualizations are currently supported for the Earth and Moon. {colliding_body} is not supported."
)
impacts_on_colliding_body = impacts.apply_mask(
pc.equal(impacts.collision_object.code, colliding_body)
)
impacting_variants_body = impacting_variants.apply_mask(
pc.is_in(impacting_variants.orbit_id, impacts_on_colliding_body.variant_id)
)
for impact in impacts_on_colliding_body:
post_impact_mask = pc.and_(
pc.is_in(impacting_variants_body.orbit_id, impact.variant_id),
pc.greater_equal(
impacting_variants_body.coordinates.time.mjd(),
impact.coordinates.time.rescale(
impacting_variants_body.coordinates.time.scale
).mjd()[0],
),
)
impacting_variants_body_correct = impacting_variants_body.apply_mask(
pc.invert(post_impact_mask)
)
impacting_variants_body_incorrect = impacting_variants_body.apply_mask(
post_impact_mask
)
if len(impacting_variants_body_incorrect) > 0:
collision_coordinates = impact.collision_coordinates
collision_coordinates = qv.concatenate(
[
collision_coordinates
for _ in range(len(impacting_variants_body_incorrect))
]
)
# Hack: In some cases, due to the time step of the propagation, the variants may already
# be inside the sphere of the colliding body when the impact is detected. In these cases, the variants'
# distance from the bodycenter is set its radius.
collision_coordinates = collision_coordinates.set_column(
"rho", pa.repeat(radius / KM_P_AU, len(collision_coordinates))
)
# Override the time of the impact coordinates to the time of the propagated variants beyond the impact time, we do
# this so we can then calculate position of these locations on the surface of the Earth as the Earth rotates.
collision_coordinates = collision_coordinates.set_column(
"time", impacting_variants_body_incorrect.coordinates.time
)
geocentric_coordinates = transform_coordinates(
collision_coordinates,
representation_out=CartesianCoordinates,
frame_out="ecliptic",
origin_out=OriginCodes.EARTH,
)
impacting_variants_body_incorrect = (
impacting_variants_body_incorrect.set_column(
"coordinates",
geocentric_coordinates,
)
)
impacting_variants_body = qv.concatenate(
[impacting_variants_body_correct, impacting_variants_body_incorrect]
)
prepared_variants[colliding_body] = impacting_variants_body
return prepared_variants
[docs]
def generate_impact_visualization_data(
orbit: Orbits,
variant_orbits: VariantOrbits,
impacts: CollisionEvent,
propagator: Propagator,
time_step: float = 5,
time_range: float = 60,
max_processes: Optional[int] = None,
) -> Tuple[Timestamp, Orbits, dict[str, Orbits]]:
"""
Generates the data for the impact visualization animation. The user should be careful
to only send in collision events that correspond to an impact with a planetary body or moon.
Non-impacting collisions such as close approaches have not been tested for this function.
CollisionEvents visualizations are currently supported for the Earth and Moon.
Parameters
----------
orbit: Orbits
The nominal best-fit orbit to propagate.
variant_orbits: VariantOrbits
The variants to propagate.
impacts: CollisionEvent
The impacts detected within the variants.
propagator: Propagator
The propagator to use to propagate the orbit.
time_step: float
The time step to use for the propagation.
time_range: float
The time range to use for the propagation.
max_processes: Optional[int]
The maximum number of processes to use for the propagation.
Returns
-------
Tuple[Timestamp, Orbits, dict[str, Orbits]]
The propagation times, the propagated nominal best-fit orbit and the propagated variants.
"""
if pc.any(
pc.invert(
pc.or_(
pc.equal(impacts.collision_object.code, "EARTH"),
pc.equal(impacts.collision_object.code, "MOON"),
)
)
).as_py():
raise ValueError(
"CollisionEvents visualizations are currently supported for the Earth and Moon."
)
# Calculate the range of impact times
impact_times = impacts.coordinates.time.mjd().to_numpy(zero_copy_only=False)
first_impact_time = np.min(impact_times)
last_impact_time = np.max(impact_times)
# Create propagation times around the range of impact times
mjds = np.arange(
first_impact_time - time_range / 60 / 24,
last_impact_time + time_range / 60 / 24 + time_step / 60 / 24,
time_step / 60 / 24,
)
mjds = mjds - np.mod(mjds, time_step / 60 / 24)
propagation_times = Timestamp.from_mjd(mjds, scale=impacts.coordinates.time.scale)
# Propagate the variants to the propagation times
propagated_variants = propagator.propagate_orbits(
Orbits.from_kwargs(
orbit_id=variant_orbits.variant_id,
object_id=variant_orbits.object_id,
coordinates=variant_orbits.coordinates,
),
propagation_times,
max_processes=max_processes,
)
# Transform the propagated variants to geocentric frame
propagated_variants = propagated_variants.set_column(
"coordinates",
transform_coordinates(
propagated_variants.coordinates,
representation_out=CartesianCoordinates,
frame_out="ecliptic",
origin_out=OriginCodes.EARTH,
),
)
# Propagate the nominal best-fit orbit to the propagation times
propagated_orbit = propagator.propagate_orbits(
orbit, propagation_times, max_processes=max_processes
)
# Transform the best-fit orbit to geocentric frame
propagated_orbit = propagated_orbit.set_column(
"coordinates",
transform_coordinates(
propagated_orbit.coordinates,
representation_out=CartesianCoordinates,
frame_out="ecliptic",
origin_out=OriginCodes.EARTH,
),
)
propagated_variants = prepare_propagated_variants(propagated_variants, impacts)
for k, v in propagated_variants.items():
propagated_variants[k] = v.sort_by(
["coordinates.time.days", "coordinates.time.nanos"]
)
return propagation_times, propagated_orbit, propagated_variants
[docs]
def create_sphere(radius, offset=None):
"""
Create a set of points that form a sphere.
Parameters
----------
radius: float
The radius of the sphere.
offset: array-like, optional
The offset of the sphere from the origin.
"""
if offset is None:
offset = np.array([0, 0, 0])
phi = np.linspace(0, np.pi, 50)
theta = np.linspace(0, 2 * np.pi, 50)
phi, theta = np.meshgrid(phi, theta)
x = radius * np.sin(phi) * np.cos(theta)
y = radius * np.sin(phi) * np.sin(theta)
z = radius * np.cos(phi)
x += offset[0]
y += offset[1]
z += offset[2]
return x, y, z
[docs]
def add_earth(
time,
coastlines: bool = True,
origin: OriginCodes = OriginCodes.EARTH,
frame: str = "ecliptic",
show: bool = True,
) -> Tuple[go.Surface, List[go.Scatter3d]]:
"""
Add the Earth to the plot.
Parameters
----------
time: Timestamp
The time of the plot.
coastlines: bool, optional
Whether to add the coastlines to the plot.
origin: Origin, optional
The origin of the plot.
frame: str, optional
The frame of the plot.
show: bool, optional
Whether to show the Earth by default.
Returns
-------
Tuple[go.Surface, List[go.Scatter3d]]
The Earth rendered as a sphere and the traces for the coastlines.
"""
earth_state = get_perturber_state(
OriginCodes.EARTH,
time,
frame=frame,
origin=origin,
)
x, y, z = create_sphere(EARTH_RADIUS_KM * 0.999, offset=earth_state.r[0] * KM_P_AU)
surface_traces = []
if coastlines:
world = gpd.read_file(str(Coastlines))
for idx, row in world.iterrows():
# Get the polygon coordinates
if row.geometry.geom_type == "Polygon":
coords = np.array(row.geometry.exterior.coords)
elif row.geometry.geom_type == "MultiPolygon":
coords = np.array(row.geometry.geoms[0].exterior.coords)
elif row.geometry.geom_type == "LineString":
coords = np.array(row.geometry.coords)
elif row.geometry.geom_type == "MultiLineString":
coords = np.array(row.geometry.geoms[0].coords)
coords = SphericalCoordinates.from_kwargs(
rho=pa.repeat(EARTH_RADIUS_KM / KM_P_AU, len(coords)),
lon=coords[:, 0],
lat=coords[:, 1],
time=Timestamp.from_kwargs(
days=pa.repeat(time.days[0], len(coords)),
nanos=pa.repeat(time.nanos[0], len(coords)),
scale=time.scale,
),
frame="itrf93",
origin=Origin.from_kwargs(
code=pa.repeat(OriginCodes.EARTH.name, len(coords))
),
)
coords = transform_coordinates(
coords,
representation_out=CartesianCoordinates,
frame_out=frame,
origin_out=origin,
)
surface_traces.append(
go.Scatter3d(
x=coords.x.to_numpy(zero_copy_only=False) * KM_P_AU,
y=coords.y.to_numpy(zero_copy_only=False) * KM_P_AU,
z=coords.z.to_numpy(zero_copy_only=False) * KM_P_AU,
mode="lines",
line=dict(color="white", width=2),
showlegend=False,
visible=show,
legendgroup="Earth",
)
)
earth_surface = go.Surface(
x=x,
y=y,
z=z,
opacity=1,
colorscale=[[0, "#015294"], [1, "#015294"]],
showscale=False,
name="Earth",
legendgroup="Earth",
showlegend=True,
visible=show,
)
return earth_surface, surface_traces
[docs]
def add_moon(
time: Timestamp,
origin: OriginCodes = OriginCodes.EARTH,
frame: Literal["ecliptic", "equatorial", "itrf93"] = "ecliptic",
show: bool = True,
) -> go.Surface:
"""
Add the Moon to the plot.
Parameters
----------
time: Timestamp
The time of the snapshot.
origin: OriginCodes
The origin of the plot.
frame: Literal["ecliptic", "equatorial", "itrf93"]
The frame of the plot.
Returns
-------
go.Surface
The Moon rendered as a sphere.
"""
lunar_state = get_perturber_state(
OriginCodes.MOON,
time,
frame=frame,
origin=origin,
)
x, y, z = create_sphere(MOON_RADIUS_KM, offset=lunar_state.r[0] * KM_P_AU)
return go.Surface(
x=x,
y=y,
z=z,
opacity=1.0,
colorscale=[[0, "#A9A9A9"], [1, "#A9A9A9"]],
showscale=False,
name="Moon",
showlegend=True,
visible=show,
)
[docs]
def plot_impact_simulation(
propagation_times: Timestamp,
propagated_best_fit_orbit: Orbits,
propagated_variants: dict[str, Orbits],
impacts: CollisionEvent,
grid: bool = True,
title: str = None,
logo: bool = True,
show_impacting: bool = True,
show_non_impacting: bool = True,
show_best_fit: bool = True,
show_earth: bool = True,
show_moon: bool = True,
sample_impactors: Optional[float] = None,
sample_non_impactors: Optional[float] = None,
height: Optional[int] = None,
width: Optional[int] = None,
) -> go.Figure:
"""
Plot the impact simulation.
Parameters
----------
propagated_best_fit_orbit: Orbits
The propagated best-fit orbit.
propagated_variants: Orbits
The propagated variants.
impacts: CollisionEvent
The impacts detected within the variants.
grid: bool, optional
Whether to add the grid to the plot.
title: str, optional
The title of the plot.
logo: bool, optional
Whether to add the Asteroid Institute logo to the plot.
show_impacting: bool, optional
Whether to show the impacting variants.
show_non_impacting: bool, optional
Whether to show the non-impacting variants.
show_best_fit: bool, optional
Whether to show the best-fit orbit.
show_earth: bool, optional
Whether to show the Earth.
show_moon: bool, optional
Whether to show the Moon.
sample_impactors: Optional[float], optional
Randomly sample the impactors for plotting. Should be between 0 and 1.
sample_non_impactors: Optional[float], optional
Randomly sample the non-impactors for plotting. Should be between 0 and 1.
height: int, optional
The height of the plot.
width: int, optional
The width of the plot.
Returns
-------
go.Figure
The impact simulation plot.
"""
propagation_times_isot = propagation_times.to_astropy().isot
num_variants = 0
impact_count = {}
sampled_variants = {}
for k, v in propagated_variants.items():
num_variants += len(v.orbit_id.unique())
if k == "Non-Impacting":
if sample_non_impactors is not None:
orbit_ids = v.orbit_id.unique()
numpy_orbit_ids = orbit_ids.to_numpy(
zero_copy_only=False
) # Convert to NumPy
if len(numpy_orbit_ids) == 0:
orbit_ids_sample = numpy_orbit_ids # Already an empty numpy array with correct dtype
logger.info("No non-impacting variants available to sample.")
else:
sample_size = np.ceil(
len(numpy_orbit_ids) * sample_non_impactors
).astype(int)
sample_size = min(
sample_size, len(numpy_orbit_ids)
) # Ensure sample_size <= population
orbit_ids_sample = np.random.choice(
numpy_orbit_ids,
sample_size,
replace=False,
)
logger.info(
f"Sampled {len(orbit_ids_sample)} non-impacting variants out of {len(numpy_orbit_ids)}"
)
# Create Arrow array with explicit type
arrow_orbit_ids_sample = pa.array(
orbit_ids_sample, type=v.orbit_id.type
)
sampled_variants[k] = v.__class__.from_pyarrow(
v.apply_mask(
pc.is_in(v.orbit_id, arrow_orbit_ids_sample)
).table.combine_chunks()
)
else:
sampled_variants[k] = v
if k != "Non-Impacting":
impact_count[k] = 0
if sample_impactors is not None:
orbit_ids = v.orbit_id.unique()
numpy_orbit_ids = orbit_ids.to_numpy(
zero_copy_only=False
) # Convert to NumPy
if len(numpy_orbit_ids) == 0:
orbit_ids_sample = numpy_orbit_ids # Empty array with correct dtype
logger.info(f"No impacting variants for '{k}' available to sample.")
else:
sample_size = np.ceil(
len(numpy_orbit_ids) * sample_impactors
).astype(int)
sample_size = min(
sample_size, len(numpy_orbit_ids)
) # Ensure sample_size <= population
orbit_ids_sample = np.random.choice(
numpy_orbit_ids,
sample_size,
replace=False,
)
logger.info(
f"Sampled {len(orbit_ids_sample)} impacting variants for '{k}' out of {len(numpy_orbit_ids)}"
)
# Create Arrow array with explicit type
arrow_orbit_ids_sample = pa.array(
orbit_ids_sample, type=v.orbit_id.type
)
sampled_variants[k] = v.__class__.from_pyarrow(
v.apply_mask(
pc.is_in(v.orbit_id, arrow_orbit_ids_sample)
).table.combine_chunks()
)
else:
sampled_variants[k] = v
all_potential_impactor_ids_from_impacts_table = set(
impacts.variant_id.unique().to_pylist()
)
if title is None:
prefix = ""
else:
prefix = f"{title}<br>"
# Build the individual frames for the animation
frames = []
for i, time in enumerate(propagation_times):
# 1. Get all impacts up to the current time
all_impacts_up_to_current_time = impacts.apply_mask(
pc.less_equal(impacts.coordinates.time.mjd(), time.mjd()[0])
)
# 2. Get the set of unique variant IDs that have impacted *anything* up to current time
current_frame_total_unique_impacted_ids_set = set(
all_impacts_up_to_current_time.variant_id.unique().to_pylist()
)
# 3. Update impact_count for each body (for title text)
for (
body_key
) in (
impact_count.keys()
): # These are "EARTH", "MOON", etc. as initialized earlier
impacts_on_this_body_up_to_current_time = (
all_impacts_up_to_current_time.apply_mask(
pc.equal(
all_impacts_up_to_current_time.collision_object.code, body_key
)
)
)
impact_count[body_key] = len(
impacts_on_this_body_up_to_current_time.variant_id.unique()
)
# Create the data for the frame
data = []
for k, v in sampled_variants.items():
if k == "Non-Impacting" and not show_non_impacting:
continue
if k == "Impacting" and not show_impacting:
continue
v_at_time = v.apply_mask(
pc.and_(
pc.equal(v.coordinates.time.days, time.days[0]),
pc.and_(
pc.less_equal(
v.coordinates.time.nanos, time.nanos[0].as_py() + 100000
),
pc.greater_equal(
v.coordinates.time.nanos, time.nanos[0].as_py() - 100000
),
),
)
)
if k == "Non-Impacting":
color = "#5685C3"
size = 1
name = k
else:
color = "red"
size = 2
name = f"{k.lower().capitalize()} Impacting"
x = v_at_time.coordinates.x.to_numpy(zero_copy_only=False) * KM_P_AU
y = v_at_time.coordinates.y.to_numpy(zero_copy_only=False) * KM_P_AU
z = v_at_time.coordinates.z.to_numpy(zero_copy_only=False) * KM_P_AU
data.append(
go.Scatter3d(
x=x,
y=y,
z=z,
mode="markers",
marker=dict(
size=size,
color=color,
opacity=1,
showscale=False,
),
name=name,
visible=True,
showlegend=True,
)
)
# Add best-fit orbit
if show_best_fit:
orbit_at_time = propagated_best_fit_orbit[i].coordinates
data.append(
go.Scatter3d(
x=orbit_at_time.x.to_numpy(zero_copy_only=False) * KM_P_AU,
y=orbit_at_time.y.to_numpy(zero_copy_only=False) * KM_P_AU,
z=orbit_at_time.z.to_numpy(zero_copy_only=False) * KM_P_AU,
mode="markers",
marker=dict(
size=3,
color="#F07620",
opacity=1,
showscale=False,
),
name="Best-Fit Orbit",
visible=True,
showlegend=True,
)
)
earth_surface, surface_traces = add_earth(time, show=show_earth)
data.append(earth_surface)
data.extend(surface_traces)
data.append(add_moon(time, show=show_moon))
text = f"{prefix}Time: {propagation_times_isot[i]}"
for k in impact_count:
text += f"<br>{k.lower().capitalize()} Impacts: {impact_count[k]} of {num_variants} Variants<br>{k.lower().capitalize()} Impact Probability: {impact_count[k]/num_variants * 100:.3f}%"
frame = go.Frame(
data=data,
name=str(i),
layout=dict(
title=dict(
text=text,
x=0.01,
y=0.97,
font=dict(size=14, color="white"),
),
),
)
frames.append(frame)
# If all impacting variants (from the original impacts table) have impacted by this frame, stop.
if all_potential_impactor_ids_from_impacts_table:
if all_potential_impactor_ids_from_impacts_table.issubset(
current_frame_total_unique_impacted_ids_set
):
break
# Plot the figure
fig = go.Figure(data=frames[0].data, frames=frames, layout=frames[0].layout)
if grid:
config = dict(
showgrid=True,
zeroline=True,
visible=True,
showticklabels=True,
gridcolor="rgba(128,128,128,0.1)",
showbackground=False,
backgroundcolor="rgb(0,0,0)",
tickfont=dict(color="white"),
)
else:
config = dict(
showgrid=False,
zeroline=False,
visible=False,
showticklabels=False,
gridcolor="rgba(128,128,128,0.0)",
showbackground=False,
backgroundcolor="rgb(0,0,0)",
tickfont=dict(color="white"),
)
if logo:
images = [
dict(
source=get_logo_base64(AsteroidInstituteLogoDark),
xref="paper",
yref="paper",
x=0.96,
y=-0.15,
sizex=0.20,
sizey=0.20,
xanchor="left",
yanchor="bottom",
layer="above",
)
]
else:
images = []
fig.update_layout(
scene=dict(
aspectratio=dict(x=1, y=1, z=1),
aspectmode="data",
camera=dict(
center=dict(x=0, y=0, z=0),
eye=dict(x=1.5, y=1.5, z=1.5),
up=dict(x=0, y=0, z=1),
),
bgcolor="rgb(0,0,0)",
xaxis=dict(title=dict(text="x [km]", font=dict(color="white")), **config),
yaxis=dict(title=dict(text="y [km]", font=dict(color="white")), **config),
zaxis=dict(title=dict(text="z [km]", font=dict(color="white")), **config),
),
height=height,
width=width,
autosize=True,
margin=dict(l=7, r=7, t=10, b=7, pad=0),
paper_bgcolor="rgb(0,0,0)",
plot_bgcolor="rgb(0,0,0)",
font=dict(color="white"),
images=images,
updatemenus=[
dict(
type="buttons",
showactive=False,
x=0.03,
y=-0.03,
buttons=[
dict(
label="▶",
method="animate",
args=[
None,
{
"frame": {"duration": 50, "redraw": True},
"fromcurrent": True,
},
],
),
dict(
label="⏸",
method="animate",
args=[
[None],
{
"frame": {"duration": 0, "redraw": False},
"mode": "immediate",
"transition": {"duration": 0},
},
],
),
],
)
],
sliders=[
dict(
currentvalue=dict(prefix="Time: "),
pad=dict(t=50),
len=0.90 if not logo else 0.80,
x=0.10,
y=0.0,
font=dict(color="white", size=10),
steps=[
dict(
args=[
[str(i)],
dict(
frame=dict(duration=50, redraw=True),
mode="immediate",
),
],
label=propagation_times_isot[i],
method="animate",
)
for i in range(len(frames))
],
)
],
uirevision="constant",
)
return fig
[docs]
def plot_risk_corridor(
impacts: CollisionEvent,
title: Optional[str] = None,
logo: bool = True,
height: Optional[int] = None,
width: Optional[int] = None,
map_style: Literal[
"carto-positron", "open-street-map", "white-bg"
] = "carto-positron",
) -> go.Figure:
"""
Plot the risk corridor with toggleable globe/map views.
Points colored by time with a linear scale and animated sequence.
Parameters
----------
impacts : Impact data containing coordinates
title : str, optional
Plot title
logo : bool, optional
Whether to add the Asteroid Institute logo to the plot.
height : int, optional
The height of the plot.
width : int, optional
The width of the plot.
map_style : {"carto-positron", "open-street-map", "white-bg"}, optional
Basemap style used by Plotly map tiles. Defaults to ``"carto-positron"``
to avoid OpenStreetMap tile-rate errors in shared/public viewing contexts.
Returns
-------
go.Figure
The risk corridor plot.
"""
# Filter to only include impacts on the Earth
impacts = impacts.apply_mask(pc.equal(impacts.collision_object.code, "EARTH"))
if len(impacts) == 0:
raise ValueError(
"No Earth impacts found. Other collision objects are not supported yet."
)
# Transform impact coordinates to ITRF93 Geodetic Coordinates
geodetic_impacts = transform_coordinates(
impacts.collision_coordinates,
representation_out=GeodeticCoordinates,
frame_out="itrf93",
origin_out=OriginCodes.EARTH,
)
# Sort all data by time
times = geodetic_impacts.time.to_astropy()
time_order = np.argsort(times.mjd)
lon = geodetic_impacts.lon.to_numpy(zero_copy_only=False)[time_order]
lat = geodetic_impacts.lat.to_numpy(zero_copy_only=False)[time_order]
times = times[time_order]
# Convert times to minutes since first impact
time_nums = (times.mjd - times.mjd.min()) * 24 * 60
# Calculate center
center_lon = lon[0]
center_lat = lat[0]
plot_config = dict(
height=height,
width=width,
map=dict(
style=map_style,
center=dict(lat=center_lat, lon=center_lon),
zoom=1,
),
)
color_bar_config = dict(
tickmode="array",
tickangle=0,
orientation="v", # Vertical orientation for right side
x=1.02,
y=0.5, # Centered vertically
xanchor="left",
yanchor="middle",
thickness=25,
len=0.8,
)
# Create frames for animation
frames = []
for i in range(len(lon)):
# Create ticks for min, middle (when applicable), and max
current_max_time = time_nums[i]
if i == 0:
# First frame: only show T+0 min
tick_values = [0]
tick_labels = ["T+0 min"]
elif i == 1:
# Second frame: show T+0 min and max
tick_values = [0, current_max_time]
tick_labels = ["T+0 min", f"T+{current_max_time:.0f} min"]
else:
# Third frame and beyond: show T+0 min, middle, and max
middle_time = current_max_time / 2
tick_values = [0, middle_time, current_max_time]
tick_labels = [
"T+0 min",
f"T+{middle_time:.0f} min",
f"T+{current_max_time:.0f} min",
]
frame = go.Frame(
data=[
go.Scattermap(
lon=lon[: i + 1],
lat=lat[: i + 1],
mode="markers",
marker=dict(
size=8,
color=time_nums[: i + 1],
colorscale="Viridis",
opacity=0.8,
showscale=True,
colorbar=dict(
title=dict(
side="right",
font=dict(size=12, color="black"),
),
ticktext=tick_labels,
tickvals=tick_values,
**color_bar_config,
),
),
name="Variant Impact Locations",
hovertext=[
f"Time: {t.iso}<br>Lon: {lo:.2f}°<br>Lat: {la:.2f}°<br>+{mins:.1f} min"
for t, lo, la, mins in zip(
times[: i + 1],
lon[: i + 1],
lat[: i + 1],
time_nums[: i + 1],
)
],
)
],
name=str(i),
)
frames.append(frame)
# Create the figure with initial state
fig = go.Figure(
data=[
go.Scattermap(
lon=[np.nan],
lat=[np.nan],
mode="markers",
marker=dict(
size=8,
color=[],
colorscale="Viridis",
opacity=0.8,
showscale=True,
colorbar=dict(
title=dict(
side="right",
font=dict(size=12, color="black"),
),
ticktext=["T+0 min", f"T+{time_nums[-1]:.0f} min"],
tickvals=[0, time_nums[-1]],
**color_bar_config,
),
),
name="Variant Impact Locations",
hovertext=[
f"Time: {times[0].iso}<br>Lon: {lon[0]:.2f}°<br>Lat: {lat[0]:.2f}°<br>+{time_nums[0]:.1f} min"
],
)
],
frames=frames,
)
if title is None:
title = "Risk Corridor"
if logo:
images = [
dict(
source=get_logo_base64(AsteroidInstituteLogoLight),
xref="paper",
yref="paper",
x=0.81,
y=-0.20, # Moved logo up slightly
sizex=0.18,
sizey=0.18,
xanchor="left",
yanchor="bottom",
layer="above",
)
]
else:
images = []
# Update layout
fig.update_layout(
title=dict(
text=title,
x=0.01,
y=0.98,
font=dict(size=14, color="black"),
),
autosize=True,
margin=dict(
l=7, r=120, t=30, b=30, pad=0
), # Increased right margin for vertical colorbar
**plot_config,
updatemenus=[
dict(
type="buttons",
showactive=False,
x=0.03,
y=-0.12, # Positioned between slider and colorbar
buttons=[
dict(
label="▶",
method="animate",
args=[
None,
{
"frame": {"duration": 50, "redraw": True},
"fromcurrent": True,
},
],
),
dict(
label="⏸",
method="animate",
args=[
[None],
{
"frame": {"duration": 0, "redraw": False},
"mode": "immediate",
"transition": {"duration": 0},
},
],
),
],
)
],
images=images,
sliders=[
dict(
currentvalue=dict(prefix="Variant: "),
pad=dict(t=50),
len=0.90 if not logo else 0.75,
x=0.05,
font=dict(color="black", size=10),
steps=[
dict(
args=[
[str(i)],
dict(
frame=dict(duration=50, redraw=True),
mode="immediate",
transition=dict(duration=0),
),
],
label=str(i),
method="animate",
)
for i in range(len(frames))
],
)
],
)
return fig