[BUG] volume rendering: sigma_clip normalization

yt/data_objects/image_array.py

@@ -332,10 +332,8 @@ def write_png(
             sigma_clip = clip_ratio
 
         if sigma_clip is not None:
-            nz = out[:, :, :3][out[:, :, :3].nonzero()]
-            return write_bitmap(
-                out.swapaxes(0, 1), filename, nz.mean() + sigma_clip * nz.std()
-            )
+            clip_value = self._clipping_value(sigma_clip, im=out)
+            return write_bitmap(out.swapaxes(0, 1), filename, clip_value)
         else:
             return write_bitmap(out.swapaxes(0, 1), filename)
 
@@ -443,3 +441,11 @@ def save(self, filename, png=True, hdf5=True, dataset_name=None):
             if not filename.endswith(".h5"):
                 filename = filename + ".h5"
             self.write_hdf5(filename, dataset_name)
+
+    def _clipping_value(self, sigma_clip, im=None):
+        # return the max value to clip with given a sigma_clip value. If im
+        # is None, the current instance is used
+        if im is None:
+            im = self
+        nz = im[:, :, :3][im[:, :, :3].nonzero()]
+        return nz.mean() + sigma_clip * nz.std()

yt/data_objects/tests/test_image_array.py

@@ -135,6 +135,18 @@ def test_write_image(self):
         im_arr.write_image("with_cmap", cmap_name="hot")
         im_arr.write_image("channel_1.png", channel=1)
 
+    def test_clipping_value(self):
+        im_arr = ImageArray(dummy_image(10.0, 4))
+        clip_val1 = im_arr._clipping_value(1)
+        clip_val2 = im_arr._clipping_value(1, im=im_arr)
+        assert clip_val2 == clip_val1
+
+        clip_val3 = im_arr._clipping_value(6)
+        assert clip_val3 > clip_val2
+
+        im_arr[:] = 1.0  # std will be 0, mean will be 1, so clip value will be 1
+        assert im_arr._clipping_value(1) == 1.0
+
     def tearDown(self):
         os.chdir(self.curdir)
         # clean up

yt/visualization/volume_rendering/scene.py

@@ -349,8 +349,10 @@ def save(
             ax = fig.add_axes([0, 0, 1, 1])
             ax.set_axis_off()
             out = self._last_render
-            nz = out[:, :, :3][out[:, :, :3].nonzero()]
-            max_val = nz.mean() + sigma_clip * nz.std()
+            if sigma_clip:
+                max_val = out._clipping_value(sigma_clip)
+            else:
+                max_val = out[:, :, :3].max()
             alpha = 255 * out[:, :, 3].astype("uint8")
             out = np.clip(out[:, :, :3] / max_val, 0.0, 1.0) * 255
             out = np.concatenate([out.astype("uint8"), alpha[..., None]], axis=-1)
@@ -511,11 +513,9 @@ def _show_mpl(self, im, sigma_clip=None, dpi=100):
         ax.set_position([0, 0, 1, 1])
 
         if sigma_clip is not None:
-            nz = im[im > 0.0]
-            nim = im / (nz.mean() + sigma_clip * np.std(nz))
+            nim = im / im._clipping_value(sigma_clip)
             nim[nim > 1.0] = 1.0
             nim[nim < 0.0] = 0.0
-            del nz
         else:
             nim = im
         axim = ax.imshow(nim[:, :, :3] / nim[:, :, :3].max(), interpolation="bilinear")

yt/visualization/volume_rendering/tests/test_save_render.py

@@ -40,4 +40,8 @@ def test_save_render(self):
         sc.save(os.path.join(self.tmpdir, "clip_2.png"), sigma_clip=2.0, render=False)
         sc.save(os.path.join(self.tmpdir, "clip_4.png"), sigma_clip=4.0, render=False)
 
+        # save a different format with/without sigma clips
+        sc.save(os.path.join(self.tmpdir, "no_clip.jpg"), render=False)
+        sc.save(os.path.join(self.tmpdir, "clip_2.jpg"), sigma_clip=2, render=False)
+
         return sc