From 5d04203fa2d4af4a89d489efc8d4ab2de8d14418 Mon Sep 17 00:00:00 2001
From: Ricky O'Steen <39831871+rosteen@users.noreply.github.com>
Date: Fri, 15 Nov 2024 13:23:27 -0500
Subject: [PATCH] Collapse 3D y before calculating initial model guess for
 gaussian model (#3295)

* Collapse 3D y before calculating initial model guess for gaussian model

* Changelog

* Add test

* Add test that fails on main
---
 CHANGES.rst                                       |  2 ++
 .../default/plugins/model_fitting/initializers.py |  4 ++++
 .../plugins/model_fitting/tests/test_plugin.py    | 15 +++++++++++++--
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index ce0f9ad0d4..1749755b7e 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -89,6 +89,8 @@ Cubeviz
 
 - Fixed "spectrum at spaxel" tool so it no longer resets spectral axis zoom. [#3249]
 
+- Fixed initializing a Gaussian1D model component when ``Cube Fit`` is toggled on. [#3295]
+
 Imviz
 ^^^^^
 
diff --git a/jdaviz/configs/default/plugins/model_fitting/initializers.py b/jdaviz/configs/default/plugins/model_fitting/initializers.py
index 0ec0d1c2b4..35ba04fec7 100644
--- a/jdaviz/configs/default/plugins/model_fitting/initializers.py
+++ b/jdaviz/configs/default/plugins/model_fitting/initializers.py
@@ -188,6 +188,10 @@ def initialize(self, instance, x, y):
             The initialized model.
         """
 
+        if y.ndim == 3:
+            # For cube fitting, need to collapse before these calculations
+            y = np.nanmean(y, axis=(0, 1))
+
         # X centroid estimates the position
         centroid = np.sum(x * y) / np.sum(y)
 
diff --git a/jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py b/jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py
index 4818de1640..0190eb25d2 100644
--- a/jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py
+++ b/jdaviz/configs/default/plugins/model_fitting/tests/test_plugin.py
@@ -156,9 +156,20 @@ def test_register_cube_model(cubeviz_helper, spectrum1d_cube):
     assert test_label in cubeviz_helper.app.data_collection
 
 
+def test_initialize_gaussian_with_cube(cubeviz_helper, spectrum1d_cube_larger):
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore')
+        cubeviz_helper.load_data(spectrum1d_cube_larger)
+    modelfit_plugin = cubeviz_helper.plugins['Model Fitting']
+
+    modelfit_plugin.cube_fit = True
+    # Check that we can initialize a Gaussian1D with cube fit toggled on
+    modelfit_plugin.create_model_component('Gaussian1D', 'G')
+
+
 def test_fit_cube_no_wcs(cubeviz_helper):
-    # This is like when user do something to a cube outside of Jdaviz
-    # and then load it back into a new instance of Cubeviz for further analysis.
+    # This is like when user does something to a cube outside of Jdaviz
+    # and then loads it back into a new instance of Cubeviz for further analysis.
     sp = Spectrum1D(flux=np.ones((7, 8, 9)) * u.nJy)  # nx, ny, nz
     cubeviz_helper.load_data(sp, data_label="test_cube")
     mf = cubeviz_helper.plugins['Model Fitting']