From 69e7dd23d5deedd67ff612cc7cb39808466c0c3a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?bj=C3=B8rn=20k=C3=A5re=20hegstad?=
 <35031140+bkhegstad@users.noreply.github.com>
Date: Fri, 26 May 2023 16:30:53 +0200
Subject: [PATCH] Added argument time_unit to ensemble get_volumetric_rates
 (#232)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Added argument time_unit to ensemble get_volumetric_rates

---------

Co-authored-by: Bjørn Kåre Hegstad <bkh@statoil.com>
---
 src/fmu/ensemble/ensemble.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/fmu/ensemble/ensemble.py b/src/fmu/ensemble/ensemble.py
index 04fa5f68..515462f3 100644
--- a/src/fmu/ensemble/ensemble.py
+++ b/src/fmu/ensemble/ensemble.py
@@ -804,7 +804,7 @@ def load_smry(
             time_index = "raw"
         return self.get_df("share/results/tables/unsmry--" + time_index + ".csv")
 
-    def get_volumetric_rates(self, column_keys=None, time_index=None):
+    def get_volumetric_rates(self, column_keys=None, time_index=None, time_unit=None):
         """Compute volumetric rates from cumulative summary vectors
 
         Column names that are not referring to cumulative summary
@@ -819,9 +819,21 @@ def get_volumetric_rates(self, column_keys=None, time_index=None):
         opposed to rates coming directly from the Eclipse simulator which
         are valid backwards in time.
 
+        If time_unit is set, the rates will be scaled to represent
+        either daily, monthly or yearly rates. These will sum up to the
+        cumulative as long as you multiply with the correct number
+        of days, months or year between each consecutive date index.
+        Month lengths and leap years are correctly handled.
+
         Args:
             column_keys (str or list of str): cumulative summary vectors
             time_index (str or list of datetimes):
+            time_unit: str or None. If None, the rates returned will
+                be the difference in cumulative between each included
+                time step (where the time interval can vary arbitrarily)
+                If set to 'days', 'months' or 'years', the rates will
+                be scaled to represent a daily, monthly or yearly rate that
+                is compatible with the date index and the cumulative data.
 
         Returns:
             pd.DataFrame: analoguous to the dataframe returned by get_smry().
@@ -830,7 +842,7 @@ def get_volumetric_rates(self, column_keys=None, time_index=None):
         vol_dfs = []
         for realidx, real in self.realizations.items():
             vol_real = real.get_volumetric_rates(
-                column_keys=column_keys, time_index=time_index
+                column_keys=column_keys, time_index=time_index, time_unit=time_unit
             )
             if "DATE" not in vol_real.columns and vol_real.index.name == "DATE":
                 # This should be true, if not we might be in trouble.