From c4595f3e47cefcaf8f3689d647dbd69fe576c905 Mon Sep 17 00:00:00 2001
From: Adrian Price-Whelan <adrian.prw@gmail.com>
Date: Wed, 14 Feb 2024 17:22:39 -0500
Subject: [PATCH] if an ax is not provided, need to get fig from ax

---
 gala/potential/potential/core.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/gala/potential/potential/core.py b/gala/potential/potential/core.py
index 7d0c4846..f782948d 100644
--- a/gala/potential/potential/core.py
+++ b/gala/potential/potential/core.py
@@ -757,6 +757,8 @@ def plot_rotation_curve(self, R_grid, t=0.0, ax=None, labels=None, **plot_kwargs
 
         if ax is None:
             fig, ax = plt.subplots()
+        else:
+            fig = ax.figure
 
         if labels is not False:
             ax.set_xlabel(labels[0])