From 496261068b0778bf9b81d67e821a684ad8d6b0ae Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Tue, 20 Apr 2021 10:56:38 +0200
Subject: [PATCH] [postprocessor] Added oct_vect_getter in generic rules

---
 baseprocessor.py |  4 ++++
 postprocessor.py | 10 +++++++++-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/baseprocessor.py b/baseprocessor.py
index 081af3a..9f8b8b1 100644
--- a/baseprocessor.py
+++ b/baseprocessor.py
@@ -457,5 +457,9 @@ def vect_getter(name, i, dset):
     return dset[name][:, i]
 
 
+def oct_vect_getter(name, i, dset):
+    return dset[name][:, :, i]
+
+
 def norm_getter(name, dset):
     return np.sqrt(np.sum(dset[name] ** 2, axis=1))
diff --git a/postprocessor.py b/postprocessor.py
index 61a2882..e9d5943 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -32,7 +32,14 @@ import pspec_new
 
 
 from units import U
-from baseprocessor import HDF5Container, Rule, norm_getter, simple_getter, vect_getter
+from baseprocessor import (
+    HDF5Container,
+    Rule,
+    norm_getter,
+    simple_getter,
+    vect_getter,
+    oct_vect_getter,
+)
 
 
 # Getters
@@ -1706,6 +1713,7 @@ class PostProcessor(HDF5Container):
                         field + dir,
                         partial(vect_getter, field, i),
                         self.unit_key[field],
+                        oct_getter=partial(oct_vect_getter, field, i),
                     )
 
                 # Radial