Skip to content

Commit

Permalink
Make impact plot style and spacing consistent
Browse files Browse the repository at this point in the history
  • Loading branch information
moeyensj committed Feb 19, 2025
1 parent ed8761a commit 1f73df5
Show file tree
Hide file tree
Showing 2 changed files with 136 additions and 64 deletions.
118 changes: 79 additions & 39 deletions examples/2024_yr4_impact_risk.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -17,27 +17,27 @@
"Before we can start, you will need to install a couple packages.\n",
"\n",
"```bash\n",
"pip install adam-core adam-assist\n",
"pip install adam-core[plots] adam-assist\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": 13,
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" orbit_id object_id coordinates.x coordinates.y coordinates.z \\\n",
"0 00000 (2024 YR4) -0.554305 1.039765 -0.031525 \n",
"0 00000 (2024 YR4) -0.59325 1.039129 -0.033844 \n",
"\n",
" coordinates.vx coordinates.vy coordinates.vz coordinates.time.days \\\n",
"0 -0.019574 -0.000133 -0.001166 60693 \n",
"0 -0.019371 -0.0005 -0.001154 60695 \n",
"\n",
" coordinates.time.nanos coordinates.covariance.values \\\n",
"0 0 [2.6790681662875186e-12, -4.746648923989143e-1... \n",
"0 0 [7.787687818777601e-13, -1.3878522864717694e-1... \n",
"\n",
" coordinates.origin.code \n",
"0 SUN \n"
Expand All @@ -48,22 +48,22 @@
"# Fetch the state vector from Small Bodies Database (SBDB)\n",
"from adam_core.orbits.query import query_sbdb\n",
"\n",
"yr4 = query_sbdb([\"2024 YR4\"])\n",
"orbit = query_sbdb([\"2024 YR4\"])\n",
"\n",
"print(yr4.to_dataframe())"
"print(orbit.to_dataframe())"
]
},
{
"cell_type": "code",
"execution_count": 14,
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[\n",
" 2925\n",
" 2923\n",
"]\n"
]
}
Expand All @@ -75,20 +75,22 @@
"from adam_core.time import Timestamp\n",
"approx_impact_date = Timestamp.from_iso8601([\"2032-12-22\"], scale=\"tdb\")\n",
"thirty_days_after_impact = approx_impact_date.add_days(30)\n",
"days_until_thirty_days_after_impact, _ = thirty_days_after_impact.difference(yr4.coordinates.time)\n",
"days_until_thirty_days_after_impact, _ = thirty_days_after_impact.difference(orbit.coordinates.time)\n",
"print(days_until_thirty_days_after_impact)"
]
},
{
"cell_type": "code",
"execution_count": 15,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"name": "stderr",
"output_type": "stream",
"text": [
"EarthImpacts(size=19)\n"
"/Users/moeyensj/software/anaconda3/lib/python3.11/subprocess.py:1883: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.\n",
" self.pid = _fork_exec(\n",
"2025-02-19 10:25:19,772\tINFO worker.py:1841 -- Started a local Ray instance.\n"
]
}
],
Expand All @@ -99,48 +101,86 @@
"\n",
"propagator = ASSISTPropagator()\n",
"\n",
"final_states, earth_impacts = calculate_impacts(\n",
" yr4,\n",
"variants, impacts = calculate_impacts(\n",
" orbit,\n",
" days_until_thirty_days_after_impact[0].as_py(),\n",
" propagator,\n",
" num_samples=1000,\n",
" num_samples=10000,\n",
" processes=10, # Multiprocessing speeds things up if you have the CPUs\n",
")\n",
"\n",
"print(earth_impacts)"
"print(variants, impacts)"
]
},
{
"cell_type": "code",
"execution_count": 16,
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
" orbit_id impacts variants cumulative_probability mean_impact_time.days \\\n",
"0 00000 19 1000 0.019 63588 \n",
"\n",
" mean_impact_time.nanos stddev_impact_time minimum_impact_time.days \\\n",
"0 50517116965167 0.00501 63588 \n",
"\n",
" minimum_impact_time.nanos maximum_impact_time.days \\\n",
"0 49983530294895 63588 \n",
"\n",
" maximum_impact_time.nanos \n",
"0 51346236683428 \n"
]
}
],
"outputs": [],
"source": [
"# Now we can summarize the returns, a simple ratio in our case\n",
"# with only 1 orbit considered\n",
"from adam_core.dynamics.impacts import calculate_impact_probabilities\n",
"ip = calculate_impact_probabilities(final_states, earth_impacts)\n",
"ip = calculate_impact_probabilities(variants, impacts)\n",
"print(ip.to_dataframe())"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We've recently added some visualization functionality to visualize both the risk corridor and the impact events. Let's first look at the impact corridor."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from adam_core.dynamics.plots import plot_risk_corridor\n",
"\n",
"fig = plot_risk_corridor(impacts, title=\"Risk Corridor for 2024 YR4\")\n",
"fig.show()"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"from adam_core.dynamics.plots import generate_impact_visualization_data\n",
"\n",
"propagated_best_fit_orbit, propagated_variants = generate_impact_visualization_data(\n",
" orbit,\n",
" variants,\n",
" impacts,\n",
" propagator,\n",
" time_step=5,\n",
" time_range=120,\n",
" max_processes=None\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from adam_core.dynamics.plots import plot_impact_simulation\n",
"\n",
"fig = plot_impact_simulation(\n",
" propagated_best_fit_orbit, \n",
" propagated_variants, \n",
" impacts, \n",
" grid=True, \n",
" title=\"2024 YR4 Impact Simulation\",\n",
")\n",
"fig.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
Expand Down Expand Up @@ -170,7 +210,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.7"
"version": "3.11.5"
}
},
"nbformat": 4,
Expand Down
82 changes: 57 additions & 25 deletions src/adam_core/dynamics/plots.py
Original file line number Diff line number Diff line change
Expand Up @@ -633,7 +633,9 @@ def plot_impact_simulation(
return fig


def plot_risk_corridor(impacts: EarthImpacts, title: Optional[str] = None):
def plot_risk_corridor(
impacts: EarthImpacts, title: Optional[str] = None, logo: bool = True
) -> go.Figure:
"""
Plot the risk corridor with toggleable globe/map views.
Points colored by time with a linear scale and animated sequence.
Expand All @@ -643,6 +645,13 @@ def plot_risk_corridor(impacts: EarthImpacts, title: Optional[str] = None):
impacts : Impact data containing coordinates
title : str, optional
Plot title
logo : bool, optional
Whether to add the Asteroid Institute logo to the plot.
Returns
-------
go.Figure
The risk corridor plot.
"""

# Sort all data by time
Expand Down Expand Up @@ -770,23 +779,45 @@ def plot_risk_corridor(impacts: EarthImpacts, title: Optional[str] = None):
if title is None:
title = "Risk Corridor"

if logo:
images = [
dict(
source=str(AsteroidInstituteLogoLight),
xref="paper",
yref="paper",
x=1.03,
y=-0.30,
sizex=0.19,
sizey=0.19,
xanchor="left",
yanchor="bottom",
layer="above",
)
]
else:
images = []

# Update layout
fig.update_layout(
title=dict(
text=title, xanchor="left", yanchor="top", font=dict(size=16, color="black")
text=title,
x=0.01,
y=0.98,
font=dict(size=14, color="black"),
),
autosize=False,
height=720,
width=1200,
showlegend=True,
paper_bgcolor="white",
margin=dict(l=5, r=5, t=30, b=5),
margin=dict(l=7, r=7, t=30, b=7),
**plot_config,
updatemenus=[
dict(
type="buttons",
showactive=False,
x=0.05,
y=-0.3,
x=0.03,
y=-0.03,
buttons=[
dict(
label="▶",
Expand Down Expand Up @@ -814,30 +845,31 @@ def plot_risk_corridor(impacts: EarthImpacts, title: Optional[str] = None):
],
)
],
# Add slider for manual control
images=images,
sliders=[
{
"currentvalue": {"prefix": "Variant: "},
"pad": {"t": 50},
"len": 0.9,
"x": 0.1,
"y": -0.2,
"steps": [
{
"args": [
dict(
currentvalue=dict(prefix="Variant: "),
pad=dict(t=50),
len=0.85,
x=0.10,
y=0.07,
font=dict(color="black", size=10),
steps=[
dict(
args=[
[str(i)],
{
"frame": {"duration": 50, "redraw": True},
"mode": "immediate",
"transition": {"duration": 0},
},
dict(
frame=dict(duration=50, redraw=True),
mode="immediate",
transition=dict(duration=0),
),
],
"label": str(i),
"method": "animate",
}
label=str(i),
method="animate",
)
for i in range(len(frames))
],
}
)
],
)

Expand Down Expand Up @@ -886,7 +918,7 @@ def plot_risk_corridor_globe(
width=1200,
showlegend=True,
paper_bgcolor="white",
margin=dict(l=5, r=5, t=30, b=5),
margin=dict(l=7, r=7, t=10, b=7),
)

geo_config = dict(
Expand Down

0 comments on commit 1f73df5

Please sign in to comment.