- [ENH] Added check when pushing to LE

- [ENH] Better handling of data_vars (the keys of a dataset) for structuredData

Author: Leguark

subsurface/interfaces/liquid_earth/rest_client.py

@@ -87,7 +87,10 @@ def add_new_project(self, project_name, extent, project_id=None):
 
     def add_data_to_project(self, project_id: str, data_name: str, data_type: DataTypes,
                             header: json, body: bytearray):
-
+        # Check if _ is in data_name
+        if "_" in project_id:
+            raise ValueError("project_id should not contain _")
+        
         self._post_update_meta_data(project_id, data_name, data_type)
         self._put_file_in_project(project_id, data_name + ".le", data_type, body)
         self._put_file_in_project(project_id, data_name + ".json", data_type,  json.dumps(header))

subsurface/structs/base_structures/structured_data.py

@@ -1,17 +1,16 @@
 from dataclasses import dataclass
-from typing import Dict, List
+from typing import Dict, List, Tuple
 
 import numpy as np
 import xarray as xr
 
-
 __all__ = ['StructuredData', ]
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=False)
 class StructuredData:
     data: xr.Dataset
-    data_array_name: str = "data_array"
+    _data_array_name: str = "data_array"
 
     """Primary structure definition for structured data
 
@@ -31,6 +30,17 @@ class StructuredData:
         data (xarray.Dataset)
     """
 
+    @property
+    def data_array_name(self):
+        data_var_list = list(self.data.data_vars.keys())
+        if self._data_array_name not in data_var_list:
+            raise ValueError("data_array_name not found in data_vars: {}".format(data_var_list))
+        return self._data_array_name
+
+    @data_array_name.setter
+    def data_array_name(self, data_array_name: str):
+        self._data_array_name = data_array_name
+
     @classmethod
     def from_numpy(cls, array: np.ndarray, coords: dict = None, data_array_name: str = "data_array",
                    dim_names: List[str] = None):
@@ -58,20 +68,27 @@ def values(self):
         return self.data[self.data_array_name].values
 
     @property
-    def default_dataset(self):
+    def default_data_array(self):
         return self.data[self.data_array_name]
 
-    def to_binary(self, order='F'):
-        bytearray_le = self._to_bytearray(order)
-        header = self._set_binary_header()
+    def default_data_array_to_binary(self, order='F'):
+        bytearray_le = self._to_bytearray(self.default_data_array, order)
+        header = self._set_binary_header(self.default_data_array)
 
         return bytearray_le, header
-
-    def _set_binary_header(self):
-        header = {"data_shape": self.values.shape}
+    
+    def to_binary(self, data_array: xr.DataArray, order: str = 'F') -> Tuple[bytes, Dict]:
+        bytearray_le = self._to_bytearray(data_array, order)
+        header = self._set_binary_header(data_array)
+        return bytearray_le, header
+        
+    @staticmethod
+    def _set_binary_header(data_array: xr.DataArray) -> Dict:
+        header = {"data_shape": data_array.shape}
         return header
 
-    def _to_bytearray(self, order):
-        data = self.values.astype('float32').tobytes(order)
+    @staticmethod
+    def _to_bytearray(data_array: xr.DataArray, order: str) -> bytes:
+        data = data_array.values.astype('float32').tobytes(order)
         bytearray_le = data
         return bytearray_le

subsurface/writer/to_binary.py

@@ -2,7 +2,7 @@
 
 
 def base_structs_to_binary_file(path, base_struct, order='F'):
-    bytearray_le, header = base_struct.to_binary(order=order)
+    bytearray_le, header = base_struct.default_data_array_to_binary(order=order)
     with open(path+'.json', 'w') as outfile:
         json.dump(header, outfile)
 

tests/test_interfaces/test_liquid_earth.py

@@ -35,7 +35,7 @@ def test_upload_data_to_new_project(self, liquid_earth_client, data_path):
 
         unstruc = self.test_dataset2(data_path)
 
-        body, header = unstruc.to_binary()
+        body, header = unstruc.default_data_array_to_binary()
 
         liquid_earth_client.add_data_to_project(
             project_id="museum",
@@ -62,7 +62,7 @@ def test_add_new_project(self, liquid_earth_client):
     def test_add_data_to_project(self, liquid_earth_client, data_path):
         us = read_unstruct(data_path + '/interpolator_meshes.nc')
         print(us.extent)
-        body, header = us.to_binary()
+        body, header = us.default_data_array_to_binary()
 
         liquid_earth_client.add_data_to_project(
             project_id="52f88baa-c84c-4082-bbba-869ef3819004",
@@ -112,7 +112,7 @@ def test_dataset2(self, data_path):
     def test_put_file_in_project(self, liquid_earth_client, data_path):
         us = read_unstruct(data_path + '/interpolator_meshes.nc')
         print(us.extent)
-        body, header = us.to_binary()
+        body, header = us.default_data_array_to_binary()
 
         liquid_earth_client._put_file_in_project(
             project_id="52f88baa-c84c-4082-bbba-869ef3819004",

tests/test_interfaces/test_to_binary.py

@@ -30,7 +30,7 @@ def wells(data_path):
 
 
 def test_wells_to_binary(wells):
-    bytearray_le, header = wells.to_binary()
+    bytearray_le, header = wells.default_data_array_to_binary()
     print(header)
 
     with open('well_f.json', 'w') as outfile:
@@ -50,7 +50,7 @@ def test_profile_to_binary(data_path):
 
     cross = imageio.imread(data_path + '/profiles/Profil1_cropped.png')
     struct = StructuredData.from_numpy(np.array(cross))
-    texture_binary, texture_header = struct.to_binary()
+    texture_binary, texture_header = struct.default_data_array_to_binary()
 
     origin = [traces.loc[0, 'geometry'].xy[0][0],
               traces.loc[0, 'geometry'].xy[1][0],
@@ -78,7 +78,7 @@ def test_profile_to_binary(data_path):
     import pandas as pd
 
     unstruct = UnstructuredData.from_array(v, e, vertex_attr=pd.DataFrame(uv, columns=['u', 'v']))
-    mesh_binary, mesh_header = unstruct.to_binary()
+    mesh_binary, mesh_header = unstruct.default_data_array_to_binary()
 
     with open('mesh_uv.json', 'w') as outfile:
         import json

tests/test_io/test_segy_reader.py

@@ -51,7 +51,7 @@ def get_images() -> List[str]:
 def test_converted_to_structured_data(get_structured_data):
     for x in get_structured_data:
         assert isinstance(x, StructuredData)
-        x.default_dataset.plot()
+        x.default_data_array.plot()
         plt.show(block=False)