From fe5c73cb0ed4ccd89dd3ff22180c6b7c85bd4d62 Mon Sep 17 00:00:00 2001
From: Lua-sketch <70654240+Lua-sketch@users.noreply.github.com>
Date: Mon, 27 May 2024 17:08:51 -0400
Subject: [PATCH] Update plot.py

modify plot_contour to include a "collapse" parameter - allows users to decide which direction to "slice" along if dealing with 3d data (the function previously only allowed taking slices in x, making it hard to plot the contours in any view that wasn't yz).
---
 astrodendro/plot.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/astrodendro/plot.py b/astrodendro/plot.py
index f675d44..65fb94f 100644
--- a/astrodendro/plot.py
+++ b/astrodendro/plot.py
@@ -110,7 +110,7 @@ def plot_tree(self, ax, structure=None, subtree=True, autoscale=True, **kwargs):
             ax.margins(0.05)
             ax.autoscale_view(True, True, True)
 
-    def plot_contour(self, ax, structure=None, subtree=True, slice=None, **kwargs):
+    def plot_contour(self, ax, structure=None, subtree=True, slice=None, collapse=None, **kwargs):
         """
         Plot a contour outlining all pixels in the dendrogram, or a specific.
         structure.
@@ -129,6 +129,10 @@ def plot_contour(self, ax, structure=None, subtree=True, slice=None, **kwargs):
             If dealing with a 3-d cube, the slice at which to plot the contour.
             If not set, the slice containing the peak of the structure will be
             shown
+        collapse: int, optional
+            If dealing with a 3-d cube, the dimension which will be "collapsed" when plotting.
+            I.e., if slice=85 and collapse=0, contours will be plotted by taking the 85th slice along x, 
+            "collapsing" the x direction.
 
         Notes
         -----
@@ -150,7 +154,12 @@ def plot_contour(self, ax, structure=None, subtree=True, slice=None, **kwargs):
                 if slice is None:
                     peak_index = structure.get_peak(subtree=subtree)
                     slice = peak_index[0][0]
-                mask = mask[slice, :, :]
+                if collapse is None or collapse==0:
+                    mask = mask[slice, :, :]
+                elif collapse==1:
+                    mask=mask[:, slice, :]
+                elif collapse==2:
+                    mask=mask[:, :, slice]
 
         # fix a common mistake when trying to set the color of contours
         if 'color' in kwargs and 'colors' not in kwargs: