Update xarray api to output dump the results into the

Author: haihabi

pynncml/__init__.py

@@ -2,10 +2,10 @@
 from pynncml import datasets
 from pynncml import single_cml_methods as scm
 from pynncml import multiple_cmls_methods as mcm
-from pynncml import single_sml_methods as ssm
-from pynncml import metrics
+from pynncml import sml_methods as ssm
+from pynncml.cml_nn_training import metrics
 from pynncml import neural_networks
-from pynncml import training_helpers
+from pynncml import cml_nn_training
 from pynncml import simulation
 from pynncml.plot_common import change_x_axis_time_format, plot_wet_dry_detection_mark
 

pynncml/apis/xarray_processing/xarray_inference_engine.py

@@ -1,4 +1,4 @@
-from pynncml.apis.xarray_processing.xarray_inference_engine import XarrayInferenceEngine
+from pynncml.cml_methods.apis.xarray_processing.xarray_inference_engine import XarrayInferenceEngine
 from pynncml.neural_networks import DNNType
 from pynncml.single_cml_methods.wet_dry import wet_dry_network,statistics_wet_dry
 

pynncml/apis/__init__.py pynncml/cml_methods/__init__.pypynncml/apis/__init__.py renamed to pynncml/cml_methods/__init__.py

@@ -0,0 +1,48 @@
+
+from torch import nn
+import numpy as np
+import xarray as xr
+
+from pynncml.cml_methods.base_cml_method import BaseCMLProcessingMethod
+from pynncml.cml_methods.results_data_structure import CMLResultsDataStructure
+from pynncml.datasets.xarray_processing import xarray2link
+from pynncml.multiple_cmls_methods import InferMultipleCMLs
+
+
+class XarrayInferenceEngine(nn.Module):
+    def __init__(self,in_cml2rain_method:BaseCMLProcessingMethod,is_recurrent=True,is_attenuation=False, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.inference_engine = InferMultipleCMLs(in_cml2rain_method,is_recurrent,is_attenuation)
+
+
+    def forward(self, x_xarray):
+        link_set=xarray2link(x_xarray)
+        results_data= self.inference_engine(link_set)
+
+        new_var_dims = ('time', 'sublink_id', 'cml_id')
+        new_var_coords = {
+            'time': x_xarray.time,
+            'sublink_id': x_xarray.sublink_id,
+            'cml_id': x_xarray.cml_id
+        }
+        for rname in CMLResultsDataStructure.results_types_list():
+            x_xarray[rname] = xr.DataArray(
+            np.full((x_xarray.sizes['time'], x_xarray.sizes['sublink_id'], x_xarray.sizes['cml_id']), np.nan),
+            coords=new_var_coords,
+            dims=new_var_dims
+            )
+
+        for i, link in enumerate(link_set):
+            results=self.inference_engine.cml2rain.convert_output_results(results_data[i])
+            for rname in CMLResultsDataStructure.results_types_list():
+                # Create a new DataArray with the correct dimensions and coordinates
+                new_values = xr.DataArray(
+                    getattr(results,rname),
+                    coords=[x_xarray.time],
+                    dims=['time']
+                )
+                # Assign the new values to the dataset at the specified slice.
+                # This will create a new variable if it doesn't exist.
+                x_xarray[rname].loc[dict(sublink_id=link.sublink_id, cml_id=link.cml_id)] = new_values
+
+        return x_xarray

…nncml/apis/xarray_processing/__init__.py pynncml/cml_methods/apis/__init__.pypynncml/apis/xarray_processing/__init__.py renamed to pynncml/cml_methods/apis/__init__.py pynncml/apis/xarray_processing/__init__.py renamed to pynncml/cml_methods/apis/__init__.py

@@ -1,5 +1,6 @@
 from torch import nn
 
+from pynncml.cml_methods.results_data_structure import CMLResultsDataStructure
 from pynncml.datasets.alignment import AttenuationType
 
 
@@ -9,3 +10,6 @@ def __init__(self,input_data_type:AttenuationType,input_rate:int,output_rate:int
         self.input_data_type = input_data_type
         self.input_rate = input_rate
         self.output_rate = output_rate
+
+    def convert_output_results(self,output_tensor)->CMLResultsDataStructure:
+        raise NotImplemented

pynncml/cml_methods/apis/xarray_processing/__init__.py

@@ -0,0 +1,18 @@
+import numpy as np
+from dataclasses import dataclass
+
+
+@dataclass
+class CMLResultsDataStructure:
+    """
+    Attenuation data class
+    :param wet_dry_detection: np.ndarray
+    :param rain_estimation: np.ndarray
+    """
+    wet_dry_detection: np.ndarray
+    rain_estimation: np.ndarray
+
+
+    @staticmethod
+    def results_types_list():
+        return ["wet_dry_detection", "rain_estimation"]

…pis/xarray_processing/wet_dry_methods.py …pis/xarray_processing/wet_dry_methods.pypynncml/apis/xarray_processing/wet_dry_methods.py renamed to pynncml/cml_methods/apis/xarray_processing/wet_dry_methods.py pynncml/apis/xarray_processing/wet_dry_methods.py renamed to pynncml/cml_methods/apis/xarray_processing/wet_dry_methods.py

@@ -0,0 +1 @@
+from pynncml.cml_nn_training.compute_data_normalization import compute_data_normalization