CO2 Mass maps refactoring

src/xtgeoapp_grd3dmaps/aggregate/_co2_mass.py

@@ -85,22 +85,22 @@ def translate_co2data_to_property(
         date = str(co2_at_date.date)
         mass_as_grids = _convert_to_grid(co2_at_date, dimensions, triplets)
         if store_all or "total_co2" in maps:
-            mass_as_grids["mass-total"].to_file(
-                grid_out_dir + "/MASS_TOTAL_" + date + ".roff", fformat="roff"
+            mass_as_grids["MASS-TOTAL"].to_file(
+                grid_out_dir + "/CO2-MASS-TOTAL--" + date + ".roff", fformat="roff",
             )
-            total_mass_list.append(mass_as_grids["mass-total"])
+            total_mass_list.append(mass_as_grids["MASS-TOTAL"])
         if store_all or "dissolved_co2" in maps:
-            mass_as_grids["mass-aqu-phase"].to_file(
-                grid_out_dir + "/MASS_AQU_PHASE_" + date + ".roff",
+            mass_as_grids["MASS-AQU-PHASE"].to_file(
+                grid_out_dir + "/CO2-MASS-AQU-PHASE--" + date + ".roff",
                 fformat="roff",
             )
-            dissolved_mass_list.append(mass_as_grids["mass-aqu-phase"])
+            dissolved_mass_list.append(mass_as_grids["MASS-AQU-PHASE"])
         if store_all or "free_co2" in maps:
-            mass_as_grids["mass-gas-phase"].to_file(
-                grid_out_dir + "/MASS_GAS_PHASE_" + date + ".roff",
+            mass_as_grids["MASS-GAS-PHASE"].to_file(
+                grid_out_dir + "/CO2-MASS-GAS-PHASE--" + date + ".roff",
                 fformat="roff",
             )
-            free_mass_list.append(mass_as_grids["mass-gas-phase"])
+            free_mass_list.append(mass_as_grids["MASS-GAS-PHASE"])
 
     return [
         free_mass_list,
@@ -162,12 +162,12 @@ def _convert_to_grid(
     date = str(co2_at_date.date)
     for mass, name in zip(
         [co2_at_date.total_mass(), co2_at_date.aqu_phase, co2_at_date.gas_phase],
-        ["mass-total", "mass-aqu-phase", "mass-gas-phase"],
+        ["MASS-TOTAL", "MASS-AQU-PHASE", "MASS-GAS-PHASE"],
     ):
         mass_array = np.zeros(dimensions)
         for i, triplet in enumerate(triplets):
             mass_array[triplet] = mass[i]
-        mass_name = "co2-" + name + "--" + date
+        mass_name = "CO2-" + name
         grids[name] = xtgeo.grid3d.GridProperty(
             ncol=dimensions[0],
             nrow=dimensions[1],

src/xtgeoapp_grd3dmaps/aggregate/_config.py

@@ -87,6 +87,7 @@ class CO2MassSettings:
     unrst_source: str
     init_source: str
     maps: Optional[List[str]] = None
+    zones: Optional[List[str]] = None
 
     def __post_init__(self):
         pass

src/xtgeoapp_grd3dmaps/aggregate/_grid_aggregation.py

@@ -55,7 +55,7 @@ def aggregate_maps(
     )
     # Iterate filters
     results = _properties_to_maps(
-        inclusion_filters,
+        inclusion_filters,	
         props,
         weights,
         method,
@@ -299,8 +299,7 @@ def _property_to_map(
     assert weights is None or weights.shape == prop.shape
     if weights is not None:
         assert method in [AggregationMethod.MEAN, AggregationMethod.SUM]
-    data = prop[0][cols] if len(prop) == 1 else prop[cols]
-    # Small hack due to a small difference between calculating mass and other properties
+    data = prop[cols]
     weights = np.ones_like(data) if weights is None else weights[cols]
     if data.mask.any():
         invalid = data.mask

src/xtgeoapp_grd3dmaps/aggregate/grid3d_co2_mass.py

@@ -3,14 +3,15 @@
 import sys
 import tempfile
 import xtgeo
-from typing import List
+import yaml
+from typing import List, Optional, Dict, Tuple
 from xtgeoapp_grd3dmaps.aggregate import (
     _co2_mass,
     _config,
     _parser,
     grid3d_aggregate_map,
 )
-from xtgeoapp_grd3dmaps.aggregate._config import CO2MassSettings
+from xtgeoapp_grd3dmaps.aggregate._config import (CO2MassSettings,Zonation)
 from ccs_scripts.co2_containment.co2_calculation import calculate_co2
 
 PROPERTIES_TO_EXTRACT = [
@@ -41,82 +42,143 @@
 # """
 
 
-def calculate_mass_property(
-    grid_file: str,
-    co2_mass_settings: CO2MassSettings,
-    out_folder: _config.Output,
-) -> List[List[xtgeo.GridProperty]]:
+def generate_co2_mass_maps(config_) :
+
     """
-    Calculates and exports 3D CO2 mass properties from the provided grid and config files
+    Calculates and exports 2D and 3D CO2 mass properties from the provided config file
 
     Args:
-        grid_file (str): Path to EGRID-file
-        co2_mass_settings (CO2MassSettings): Settings from config file for calculation
-                                             of CO2 mass maps.
-        out_folder (str): Path to store the produced 3D GridProperties.
-
-
-    Returns:
-        List[List[xtgeo.GridProperty].
-
+        config_: Arguments in the config file
     """
-    co2_data = calculate_co2(grid_file,co2_mass_settings.unrst_source,"mass",co2_mass_settings.init_source,None)
 
+    co2_mass_settings = config_.co2_mass_settings
+    zonation = config_.zonation
+    zones = co2_mass_settings.zones
+    if zones is not None and isinstance(zones, str):
+        co2_mass_settings.zones = [zones]
+    grid_file = config_.input.grid
+    co2_data = calculate_co2(grid_file,co2_mass_settings.unrst_source,"mass",co2_mass_settings.init_source,None)
+    dates = config_.input.dates
+    if len(dates)>0:
+        co2_data.data_list = [x for x in co2_data.data_list if x.date in dates]
     out_property_list = _co2_mass.translate_co2data_to_property(
         co2_data,
         grid_file,
         co2_mass_settings,
         PROPERTIES_TO_EXTRACT,
-        out_folder.mapfolder + "/grid",
+        config_.output.mapfolder + "/grid",
     )
-    return out_property_list
-
+    config_.zonation.zranges, all_zrange = process_zonation(co2_mass_settings,grid_file,zonation)
+    if len(config_.zonation.zranges)>0:
+        config_.zonation.zproperty = None
+    if config_.computesettings.all:
+        config_.zonation.zranges.append({'all':all_zrange})
+        config_.computesettings.all = False
+        if not config_.computesettings.zone:
+            config_.computesettings.zone = True
+            config_.zonation.zranges = [zrange for zrange in config_.zonation.zranges if 'all' in zrange]
+    co2_mass_property_to_map(config_,out_property_list)
 
 def co2_mass_property_to_map(
     config_: _config.RootConfig,
-    t_prop: xtgeo.GridProperty,
+    property_list: List[xtgeo.GridProperty],
 ):
     """
-    Aggregates with SUM and writes a CO2 mass property to file using `grid3d_aggregate_map`.
-    The property is written to a temporary file while performing the
-    aggregation.
+    Aggregates with SUM and writes a list of CO2 mass property to files
+    using `grid3d_aggregate_map`.
 
     Args:
         config_: Arguments in the config file
-        t_prop: Grid property to be aggregated
+        property_list: List of Grid property objects to be aggregated
 
     """
     config_.input.properties = []
     config_.computesettings.aggregation = _config.AggregationMethod.SUM
     config_.output.aggregation_tag = False
-    _, temp_path = tempfile.mkstemp()
-    config_.input.properties.append(_config.Property(temp_path, t_prop.name, None))
-    t_prop.to_file(temp_path)
+    for props in property_list:
+        if len(props)>0 :
+            for prop in props:
+                config_.input.properties.append(_config.Property(config_.output.mapfolder+
+                                                                 "/grid/"+prop.name+"--"+
+                                                                 prop.date+".roff", None, None))
     grid3d_aggregate_map.generate_from_config(config_)
-    os.unlink(temp_path)
 
+def process_zonation(co2_mass_settings: _config.CO2MassSettings,
+                     grid_file: str,
+                     zonation: Optional[_config.Zonation]=None
+                     ) -> Tuple[List,List]:
+    """
+    Processes a zonation file, if existing, and extracts both zranges per zone
+    and the complete range in the zaxis. Otherwise, uses the grid_file.
+
+    Args:
+        co2_mass_settings: Arguments in CO2 mass settings
+        grid_file: Path to grid file
+        zonation: Arguments in zonation
+
+    Returns:
+        Tuple[List,List]
+    """
+    if zonation.zproperty is not None or len(zonation.zranges)>0:
+        if zonation.zproperty is not None:
+            if zonation.zproperty.source.split(".")[-1] in ["yml", "yaml"]:
+                zfile = read_yml_file(zonation.zproperty.source)
+                zonation.zranges = zfile['zranges']
+        if len(zonation.zranges) > 0:
+            zone_names = [list(item.keys())[0] for item in zonation.zranges]
+            zranges_limits = [list(d.values())[0] for d in zonation.zranges]
+    else:
+        grid_pf = xtgeo.grid_from_file(grid_file)
+        zranges_limits = [[1,grid_pf.nlay]]
+        zone_names = None
+    max_zvalue = max(sublist[-1] for sublist in zranges_limits)
+    min_zvalue = min(sublist[0] for sublist in zranges_limits)
+    all_zrange = [min_zvalue, max_zvalue]
+    if zone_names is not None:
+        if co2_mass_settings.zones is not None:
+            zones_to_plot = [zone for zone in co2_mass_settings.zones if zone in zone_names]
+            if len(zones_to_plot) == 0:
+                print(
+                    "The zones specified in CO2 mass settings are not part of the zonation provided \n maps will be exported for all the existing zones")
+                return zonation.zranges,all_zrange
+            else:
+                return [item for item in zonation.zranges if list(item.keys())[0] in zones_to_plot],all_zrange
+        else:
+            return zonation.zranges,all_zrange
+    else:
+        return [], all_zrange
+
+def read_yml_file(file_path: str) -> Dict[str,List]:
+    """
+    Reads a yml from a given path in file_path argument
+    """
+    with open(file_path, "r", encoding="utf8") as stream:
+        try:
+            zfile = yaml.safe_load(stream)
+        except yaml.YAMLError as exc:
+            print(exc)
+            sys.exit()
+    if "zranges" not in zfile:
+        error_text = "The yaml zone file must be in the format:\nzranges:\
+        \n    - Zone1: [1, 5]\n    - Zone2: [6, 10]\n    - Zone3: [11, 14])"
+        raise Exception(error_text)
+    return zfile
 
 def main(arguments=None):
     """
     Takes input arguments and calculates co2 mass as a property and aggregates it to a 2D map
-    at each time step, divided into different phases and locations(TODO).
+    at each time step, divided into different phases and locations.
     """
     if arguments is None:
         arguments = sys.argv[1:]
     config_ = _parser.process_arguments(arguments)
+
     if config_.input.properties:
         raise ValueError("CO2 mass computation does not take a property as input")
     if config_.co2_mass_settings is None:
         raise ValueError("CO2 mass computation needs co2_mass_settings as input")
-    out_property_list = calculate_mass_property(
-        config_.input.grid,
-        config_.co2_mass_settings,
-        config_.output,
-    )
 
-    for props in out_property_list:
-        for prop in props:
-            co2_mass_property_to_map(config_, prop)
+    generate_co2_mass_maps(config_)
 
 
 if __name__ == "__main__":

tests/yaml/config_co2_mass_zones_dates.yml

@@ -0,0 +1,24 @@
+input:
+  eclroot: tests/data/reek/REEK
+  grid: $eclroot.EGRID
+  dates:
+     - 21901113
+
+co2_mass_settings:
+  unrst_source: $eclroot.UNRST
+  init_source: $eclroot.INIT
+  maps: "dissolved_co2"
+  zones: ["UPPER","ZERO"]
+
+zonation:
+  zranges:
+    - UPPER: [1, 6]
+    - LOWER: [8, 14]
+    - ZERO: [15, 15]
+
+output:
+  mapfolder: tmp
+
+computesettings:
+  zone: Yes
+  all: No