From 0a618fb5fa941eee958488583541cb71643f84fb Mon Sep 17 00:00:00 2001
From: Alex Immer <aleximmer@google.com>
Date: Wed, 10 Apr 2024 09:22:39 -0700
Subject: [PATCH] Add StdProjection to compute standard deviation over given
 dimension.

PiperOrigin-RevId: 623516360
---
 connectomics/volume/decorators.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/connectomics/volume/decorators.py b/connectomics/volume/decorators.py
index 64c2126..866c01d 100644
--- a/connectomics/volume/decorators.py
+++ b/connectomics/volume/decorators.py
@@ -902,6 +902,19 @@ def __init__(self,
         context_spec=context_spec)
 
 
+@gin.register
+class StdProjection(Projection):
+  """Reduces input TensorStore along given dimension via nanstd()."""
+
+  def __init__(self,
+               projection_dim: int,
+               context_spec: Optional[MutableJsonSpec] = None):
+    super().__init__(
+        projection_fn=np.nanstd,
+        projection_dim=projection_dim,
+        context_spec=context_spec)
+
+
 @gin.register
 class SumProjection(Projection):
   """Reduces input TensorStore along given dimension via nansum()."""