Filtered mean and std from a tensor

Jorge_Garcia: <p>Hello!</p>
<p>I would like to get the mean and standard deviation from a tensor with shape (H,W) along the dimension 1 (so the output would be a pair of tensors with shape (H) or (H,1)). My problem is: I need to filter each row and only use the values selected by a given mask and the mask selects a different number of values from each row.</p>
<p>My actual approach is to generate another tensor with NaNs where I don’t care about the value and use torch.nanmean to get the mean ignoring those values but I don’t find an analogue function to get the std.</p>
<p>Example of behaviour.</p>
<pre><code class="lang-auto">&gt;&gt;&gt; input = torch.randn(2, 4)
&gt;&gt;&gt; input
tensor([[-1,4,3,8],
            [7,5,-2,-9]])
&gt;&gt;&gt; mask = input &gt;= 0
&gt;&gt;&gt; input_nans = torch.where(mask, input, torch.nan)
&gt;&gt;&gt; mean = torch.nanmean(input_nans, dim=1)
&gt;&gt;&gt; mean
tensor([5, 6])
</code></pre>
<p>This should be feasible (a mean from the values greater than 0) but there is no nanstd function to get the standard deviation ignoring nans. It’s important to note that the first value comes from the mean over 3 values while the second comes from the mean over 2 values (that is the reason I cannot simply filter to get a tensor only with the desired values, the rows length would not match).</p>
<p>¿Any suggestion?</p>

shivammehta007: <p>If you have a mask constructed you can use the mask_select or just index it.</p>
<pre><code class="lang-python">&gt;&gt;&gt; a = torch.randn(3 ,3)
&gt;&gt;&gt; a
tensor([[-1.6599, -0.0141, -0.9498],
        [ 0.5469, -0.0643, -2.0145],
        [ 0.4915, -1.7717, -1.9315]])
&gt;&gt;&gt; mask = torch.randn(3, 3) &gt; 0
&gt;&gt;&gt; mask
tensor([[False, False, False],
        [ True, False, False],
        [ True, False,  True]])
&gt;&gt;&gt; a[mask]
tensor([ 0.5469,  0.4915, -1.9315])
&gt;&gt;&gt; a[mask].mean()
tensor(-0.2977)
&gt;&gt;&gt; a[mask].std()
tensor(1.4152)
</code></pre>

jayz: <p>Maybe implement your own nanstd?</p>
<pre><code class="lang-python">def nanstd(x): 
    return torch.sqrt(torch.mean(torch.pow(x-torch.nanmean(x,dim=1).unsqueeze(-1),2)))
</code></pre>
<p>or sth similar.</p>

Jorge_Garcia: <p>That is not my desired behaviour. I’ve extended the question with an example. I hope it to be more undestable now. Thanks! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

Jorge_Garcia: <p>With that function, the output would include nans becaouse x can include nans. I’ve extended the question with an example. Thanks <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

ataxias: <p>I think that jayz 's answer needed a small modification:</p>
<pre><code class="lang-python">def nanstd(x): 
    return torch.sqrt(torch.nanmean(torch.pow(x - torch.nanmean(x, dim=-1).unsqueeze(-1), 2), dim=-1))
</code></pre>
<p>Wouldn’t this work for you?</p>

Jeremyywb: <pre><code class="lang-python">def nanstd(o,dim):
    return torch.sqrt(
                torch.nanmean(
                    torch.pow( torch.abs(o-torch.nanmean(o,dim=dim).unsqueeze(dim)),2),
                    dim=dim)
                )
nanstd(o,1)

</code></pre>
<p>this should work,</p>
<p>eg:input (bs,seq,dim)<br>
output  (bs,dim)  where you reduce seq dim to 1 to get std in seq axis without nan</p>

allenyllee: <p>Sometimes, you might want to set <code>keepdim=True</code>, here is modified version:</p>
<pre><code class="lang-auto">def nanstd(o, dim, keepdim=False):

    result = torch.sqrt(
                torch.nanmean(
                    torch.pow( torch.abs(o-torch.nanmean(o,dim=dim).unsqueeze(dim)),2),
                    dim=dim
                )
            )
    
    if keepdim:
        result = result.unsqueeze(dim)
    
    return result
</code></pre>

leowerne: <p>Thanks for pointing out keepdim.<br>
I found that keepdim instead of unsqeeze also makes it possible to pass tuples as dim.</p>
<pre data-code-wrap="cpp"><code class="lang-cpp">def nanstd(o, dim, keepdim=False):
    result = torch.sqrt(
            torch.nanmean(
                torch.pow(torch.abs(o-torch.nanmean(o,dim=dim, keepdim= True)) ,2),
                dim=dim, keepdim= keepdim
            )
        )
     return result
</code></pre>