From c5e3449ace073c1ba6b81c79c85b7d91e1b0cd74 Mon Sep 17 00:00:00 2001
From: James Trayford <j.w.trayford@port.ac.uk>
Date: Thu, 11 Jul 2024 10:49:44 +0100
Subject: [PATCH] Patches: - Fix sigcube alignment with other cubes [get rid of
 transpose] - Add percentile cut option that enables some rough feature
 isolation without full continuum subtraction - Protect against cubes w/NaNs

---
 .../configs/cubeviz/plugins/cube_listener.py  | 14 ++++++-------
 .../plugins/sonify_data/sonify_data.py        |  5 +++--
 .../plugins/sonify_data/sonify_data.vue       | 13 +++++++++++-
 jdaviz/configs/cubeviz/plugins/viewers.py     | 21 ++++++++++++++-----
 4 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/jdaviz/configs/cubeviz/plugins/cube_listener.py b/jdaviz/configs/cubeviz/plugins/cube_listener.py
index 5c1996bb50..9ce31f937e 100644
--- a/jdaviz/configs/cubeviz/plugins/cube_listener.py
+++ b/jdaviz/configs/cubeviz/plugins/cube_listener.py
@@ -74,7 +74,7 @@ def __init__(self, cube, wlens, samplerate=44100, duration=1, overlap=0.05, buff
         if self.cursig.nbytes * pow(1024,-3) > 2:
             raise Exception("Cube projected to be > 2Gb!")
             
-        self.sigcube = np.zeros((self.siglen, *self.cube.shape[1:]), dtype='int16')
+        self.sigcube = np.zeros((*self.cube.shape[:2], self.siglen), dtype='int16')
         
     def audify_cube(self, fmin=50, fmax=1500):
         """
@@ -82,22 +82,22 @@ def audify_cube(self, fmin=50, fmax=1500):
         in class attributes
         """
         lo2hi = self.wlens.argsort()[::-1]
-        for i in tqdm(range(self.cube.shape[1])):
-            for j in range(self.cube.shape[2]):
+        for i in tqdm(range(self.cube.shape[0])):
+            for j in range(self.cube.shape[1]):
                 with suppress_stderr():
-                    sig = audify_spectrum(self.cube[lo2hi,i,j], self.dur, 
+                    sig = audify_spectrum(self.cube[i,j,lo2hi], self.dur, 
                                           srate=self.srate,  
                                           fmin=fmin, fmax=fmax)
                     sig = (sig*self.maxval).astype('int16')
-                    self.sigcube[:,i,j] = sig
-        self.cursig[:] = self.sigcube[:,self.idx1,self.idx2]
+                    self.sigcube[i,j,:] = sig
+        self.cursig[:] = self.sigcube[self.idx1,self.idx2, :]
         self.newsig[:] = self.cursig[:]
 
     def player_callback(self, outdata, frames, time, status):
         cur = self.cursig
         new = self.newsig
         sdx = int(time.outputBufferDacTime*self.srate)
-        dxs = np.arange(sdx, sdx+frames).astype(int) % self.sigcube.shape[0]  
+        dxs = np.arange(sdx, sdx+frames).astype(int) % self.sigcube.shape[-1]  
         if self.cbuff:
             outdata[:,0] = (cur[dxs] * self.ofade).astype('int16')
             outdata[:,0] += (new[dxs] * self.ifade).astype('int16')
diff --git a/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.py b/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.py
index 643dfd53c5..e1403c995f 100644
--- a/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.py
+++ b/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.py
@@ -16,10 +16,11 @@ class SonifyData(PluginTemplateMixin, DatasetSelectMixin):
     ssvidx = FloatHandleEmpty(0.65).tag(sync=True)
     wavemin = FloatHandleEmpty(15800).tag(sync=True)
     wavemax = FloatHandleEmpty(16000).tag(sync=True)
-
+    pccut = IntHandleEmpty(20).tag(sync=True)
+    
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     def vue_sonify_cube(self, *args):
         viewer = self.app.get_viewer('flux-viewer')
-        viewer.get_sonified_cube(self.sample_rate, self.buffer_size, self.assidx, self.ssvidx)
+        viewer.get_sonified_cube(self.sample_rate, self.buffer_size, self.assidx, self.ssvidx, self.pccut)
diff --git a/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.vue b/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.vue
index cfd0c01aa1..a9ce2e8f6d 100644
--- a/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.vue
+++ b/jdaviz/configs/cubeviz/plugins/sonify_data/sonify_data.vue
@@ -71,7 +71,7 @@
         persistent-hint
       ></v-text-field>
     </v-row>
-        <v-row>
+    <v-row>
       <v-text-field
         ref="wavemax"
         type="number"
@@ -81,6 +81,17 @@
         persistent-hint
       ></v-text-field>
     </v-row>
+    <v-row>
+      <v-text-field
+        ref="pccut"
+        type="number"
+        label="Flux Percentile Cut"
+        v-model.number="pccut"
+        hint="The minimum flux percentile to be heard."
+        persistent-hint
+      ></v-text-field>
+    </v-row>
+   
     <v-row>
         <plugin-action-button
         @click="sonify_cube"
diff --git a/jdaviz/configs/cubeviz/plugins/viewers.py b/jdaviz/configs/cubeviz/plugins/viewers.py
index 323523a696..ed20a94dd3 100644
--- a/jdaviz/configs/cubeviz/plugins/viewers.py
+++ b/jdaviz/configs/cubeviz/plugins/viewers.py
@@ -242,18 +242,29 @@ def stop_stream(self):
     def update_cube(self, x, y):
         if not self.audified_cube or not hasattr(self.audified_cube, 'newsig') or not hasattr(self.audified_cube, 'sigcube'):
             return
-        self.audified_cube.newsig = self.audified_cube.sigcube[:, x, y]
+        self.audified_cube.newsig = self.audified_cube.sigcube[x, y, :]
         self.audified_cube.cbuff = True
 
-    def get_sonified_cube(self, sample_rate, buffer_size, assidx, ssvidx):
+    def get_sonified_cube(self, sample_rate, buffer_size, assidx, ssvidx, pccut):
         spectrum = self.active_image_layer.layer.get_object(statistic=None)
+        pc_cube = np.percentile(np.nan_to_num(spectrum.flux.value), np.clip(pccut,0,99), axis=-1)
 
-        clipped_arr = np.clip(spectrum.flux.value.T, 0, np.inf)
-        # arr = spectrum[wavemin:wavemax].flux.value.T
+        # clip zeros and remove NaNs
+        clipped_arr = np.nan_to_num(np.clip(spectrum.flux.value, 0, np.inf), copy=False)
+
+        # make a rough white-light image from the clipped array
+        whitelight = np.expand_dims(clipped_arr.sum(-1), axis=2)
+
+        # subtract any percentile cut
+        clipped_arr -= np.expand_dims(pc_cube, axis=2)
+
+        # and re-clip
+        clipped_arr = np.clip(clipped_arr, 0, np.inf)
+        
         self.audified_cube = CubeListenerData(clipped_arr ** assidx, spectrum.wavelength.value, duration=0.8,
                                   samplerate=sample_rate, buffsize=buffer_size)
         self.audified_cube.audify_cube()
-        self.audified_cube.sigcube = (self.audified_cube.sigcube * pow(clipped_arr.sum(0) / clipped_arr.sum(0).max(), ssvidx)).astype('int16')
+        self.audified_cube.sigcube = (self.audified_cube.sigcube * pow(whitelight / whitelight.max(), ssvidx)).astype('int16')
         self.stream = sd.OutputStream(samplerate=sample_rate, blocksize=buffer_size, channels=1, dtype='int16', latency='low',
                                       callback=self.audified_cube.player_callback)
         self.audified_cube.cbuff = True