ocf-data-sampler 0.1.6__tar.gz → 0.1.7__tar.gz

{ocf_data_sampler-0.1.6/ocf_data_sampler.egg-info → ocf_data_sampler-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ocf_data_sampler
-Version: 0.1.6
+Version: 0.1.7
 Summary: Sample from weather data for renewable energy prediction
 Author: James Fulton, Peter Dudfield, and the Open Climate Fix team
 Author-email: info@openclimatefix.org

ocf_data_sampler-0.1.7/ocf_data_sampler/sample/base.py

@@ -0,0 +1,75 @@
+"""
+Base class definition - abstract
+Handling of both flat and nested structures - consideration for NWP
+"""
+
+import logging
+import numpy as np
+import torch
+import xarray as xr
+
+from pathlib import Path
+from typing import Any, Dict, Optional, Union, TypeAlias
+from abc import ABC, abstractmethod
+
+
+logger = logging.getLogger(__name__)
+
+NumpySample: TypeAlias = Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]]
+NumpyBatch: TypeAlias = Dict[str, Union[np.ndarray, Dict[str, np.ndarray]]]
+TensorBatch: TypeAlias = Dict[str, Union[torch.Tensor, Dict[str, torch.Tensor]]]
+
+
+class SampleBase(ABC):
+    """ 
+    Abstract base class for all sample types 
+    Provides core data storage functionality
+    """
+
+    def __init__(self, data: Optional[Union[NumpySample, xr.Dataset]] = None):
+        """ Initialise data container """
+        logger.debug("Initialising SampleBase instance")
+        self._data = data
+
+    @abstractmethod
+    def to_numpy(self) -> NumpySample:
+        """ Convert data to a numpy array representation """
+        raise NotImplementedError
+
+    @abstractmethod
+    def plot(self, **kwargs) -> None:
+        """ Abstract method for plotting """
+        raise NotImplementedError
+
+    @abstractmethod
+    def save(self, path: Union[str, Path]) -> None:
+        """ Abstract method for saving sample data """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def load(cls, path: Union[str, Path]) -> 'SampleBase':
+        """ Abstract class method for loading sample data """
+        raise NotImplementedError
+
+
+def batch_to_tensor(batch: NumpyBatch) -> TensorBatch:
+    """
+    Moves ndarrays in a nested dict to torch tensors
+    Args:
+        batch: NumpyBatch with data in numpy arrays
+    Returns:
+        TensorBatch with data in torch tensors
+    """
+    if not batch:
+        raise ValueError("Cannot convert empty batch to tensors")
+
+    for k, v in batch.items():
+        if isinstance(v, dict):
+            batch[k] = batch_to_tensor(v)
+        elif isinstance(v, np.ndarray):
+            if v.dtype == np.bool_:
+                batch[k] = torch.tensor(v, dtype=torch.bool)
+            elif np.issubdtype(v.dtype, np.number):
+                batch[k] = torch.as_tensor(v)
+    return batch

{ocf_data_sampler-0.1.6 → ocf_data_sampler-0.1.7/ocf_data_sampler.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ocf_data_sampler
-Version: 0.1.6
+Version: 0.1.7
 Summary: Sample from weather data for renewable energy prediction
 Author: James Fulton, Peter Dudfield, and the Open Climate Fix team
 Author-email: info@openclimatefix.org

{ocf_data_sampler-0.1.6 → ocf_data_sampler-0.1.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ocf_data_sampler"
-version = "0.1.6"
+version = "0.1.7"
 license = { file = "LICENSE" }
 readme = "README.md"
 description = "Sample from weather data for renewable energy prediction"

{ocf_data_sampler-0.1.6 → ocf_data_sampler-0.1.7}/tests/test_sample/test_base.py

@@ -3,11 +3,14 @@ Base class testing - SampleBase
 """
 
 import pytest
+import torch
 import numpy as np
 
 from pathlib import Path
-from ocf_data_sampler.sample.base import SampleBase
-
+from ocf_data_sampler.sample.base import (
+    SampleBase,
+    batch_to_tensor
+)
 
 class TestSample(SampleBase):
     """
@@ -84,3 +87,61 @@ def test_sample_base_to_numpy():
     assert isinstance(numpy_data, dict)
     assert all(isinstance(value, np.ndarray) for value in numpy_data.values())
     assert np.array_equal(numpy_data['list_data'], np.array([1, 2, 3]))
+
+
+def test_batch_to_tensor_nested():
+    """ Test nested dictionary conversion """
+    batch = {
+        'outer': {
+            'inner': np.array([1, 2, 3])
+        }
+    }
+    tensor_batch = batch_to_tensor(batch)
+    
+    assert torch.equal(tensor_batch['outer']['inner'], torch.tensor([1, 2, 3]))
+
+
+def test_batch_to_tensor_mixed_types():
+    """ Test handling of mixed data types """
+    batch = {
+        'tensor_data': np.array([1, 2, 3]),
+        'string_data': 'not_a_tensor',
+        'nested': {
+            'numbers': np.array([4, 5, 6]),
+            'text': 'still_not_a_tensor'
+        }
+    }
+    tensor_batch = batch_to_tensor(batch)
+    
+    assert isinstance(tensor_batch['tensor_data'], torch.Tensor)
+    assert isinstance(tensor_batch['string_data'], str)
+    assert isinstance(tensor_batch['nested']['numbers'], torch.Tensor)
+    assert isinstance(tensor_batch['nested']['text'], str)
+
+
+def test_batch_to_tensor_different_dtypes():
+    """ Test conversion of arrays with different dtypes """
+    batch = {
+        'float_data': np.array([1.0, 2.0, 3.0], dtype=np.float32),
+        'int_data': np.array([1, 2, 3], dtype=np.int64),
+        'bool_data': np.array([True, False, True], dtype=np.bool_)
+    }
+    tensor_batch = batch_to_tensor(batch)
+    
+    assert isinstance(tensor_batch['bool_data'], torch.Tensor)
+    assert tensor_batch['float_data'].dtype == torch.float32
+    assert tensor_batch['int_data'].dtype == torch.int64
+    assert tensor_batch['bool_data'].dtype == torch.bool
+
+
+def test_batch_to_tensor_multidimensional():
+    """ Test conversion of multidimensional arrays """
+    batch = {
+        'matrix': np.array([[1, 2], [3, 4]]),
+        'tensor': np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+    }
+    tensor_batch = batch_to_tensor(batch)
+    
+    assert tensor_batch['matrix'].shape == (2, 2)
+    assert tensor_batch['tensor'].shape == (2, 2, 2)
+    assert torch.equal(tensor_batch['matrix'], torch.tensor([[1, 2], [3, 4]]))

ocf_data_sampler-0.1.6/ocf_data_sampler/sample/base.py

@@ -1,44 +0,0 @@
-"""
-Base class definition - abstract
-Handling of both flat and nested structures - consideration for NWP
-"""
-
-import logging
-import numpy as np
-
-from pathlib import Path
-from typing import Any, Dict, Optional, Union
-from abc import ABC, abstractmethod
-
-logger = logging.getLogger(__name__)
-
-class SampleBase(ABC):
-    """ 
-    Abstract base class for all sample types 
-    Provides core data storage functionality
-    """
-
-    def __init__(self):
-        """ Initialise data container """
-        logger.debug("Initialising SampleBase instance")
-
-    @abstractmethod
-    def to_numpy(self) -> Dict[str, Any]:
-        """ Convert data to a numpy array representation """
-        raise NotImplementedError
-
-    @abstractmethod
-    def plot(self, **kwargs) -> None:
-        """ Abstract method for plotting """
-        raise NotImplementedError
-
-    @abstractmethod
-    def save(self, path: Union[str, Path]) -> None:
-        """ Abstract method for saving sample data """
-        raise NotImplementedError
-
-    @classmethod
-    @abstractmethod
-    def load(cls, path: Union[str, Path]) -> 'SampleBase':
-        """ Abstract class method for loading sample data """
-        raise NotImplementedError