feat: Add orientation data loading functionality

- Add orientation schema and loading methods

- Implement orientation data validation and polarity handling

- Add comprehensive test suite for orientation loading

- Fix polarity handling in test data

Author: flohorovicic

gempy/modules/json_io/json_operations.py

@@ -12,7 +12,7 @@
 from gempy.core.data.structural_frame import StructuralFrame
 from gempy.core.data.grid import Grid
 from gempy.core.data.geo_model import GeoModel
-from .schema import SurfacePoint, GemPyModelJson
+from .schema import SurfacePoint, Orientation, GemPyModelJson
 
 
 class JsonIO:
@@ -36,11 +36,12 @@ def load_model_from_json(file_path: str) -> GeoModel:
         if not JsonIO._validate_json_schema(data):
             raise ValueError("Invalid JSON schema")
 
-        # Load surface points
+        # Load surface points and orientations
         surface_points = JsonIO._load_surface_points(data['surface_points'])
+        orientations = JsonIO._load_orientations(data['orientations'])
 
         # TODO: Load other components
-        raise NotImplementedError("Only surface points loading is implemented")
+        raise NotImplementedError("Only surface points and orientations loading is implemented")
 
     @staticmethod
     def _load_surface_points(surface_points_data: List[SurfacePoint]) -> SurfacePointsTable:
@@ -89,6 +90,69 @@ def _load_surface_points(surface_points_data: List[SurfacePoint]) -> SurfacePoin
             nugget=nugget,
             name_id_map=name_id_map
         )
+
+    @staticmethod
+    def _load_orientations(orientations_data: List[Orientation]) -> OrientationsTable:
+        """
+        Load orientations from JSON data.
+        
+        Args:
+            orientations_data (List[Orientation]): List of orientation dictionaries
+            
+        Returns:
+            OrientationsTable: A new OrientationsTable instance
+            
+        Raises:
+            ValueError: If the data is invalid or missing required fields
+        """
+        # Validate data structure
+        required_fields = {'x', 'y', 'z', 'G_x', 'G_y', 'G_z', 'id', 'nugget', 'polarity'}
+        for i, ori in enumerate(orientations_data):
+            missing_fields = required_fields - set(ori.keys())
+            if missing_fields:
+                raise ValueError(f"Missing required fields in orientation {i}: {missing_fields}")
+            
+            # Validate data types
+            if not all(isinstance(ori[field], (int, float)) for field in ['x', 'y', 'z', 'G_x', 'G_y', 'G_z', 'nugget']):
+                raise ValueError(f"Invalid data type in orientation {i}. All coordinates, gradients, and nugget must be numeric.")
+            if not isinstance(ori['id'], int):
+                raise ValueError(f"Invalid data type in orientation {i}. ID must be an integer.")
+            if not isinstance(ori['polarity'], int) or ori['polarity'] not in {-1, 1}:
+                raise ValueError(f"Invalid polarity in orientation {i}. Must be 1 (normal) or -1 (reverse).")
+        
+        # Extract coordinates and other data
+        x = np.array([ori['x'] for ori in orientations_data])
+        y = np.array([ori['y'] for ori in orientations_data])
+        z = np.array([ori['z'] for ori in orientations_data])
+        G_x = np.array([ori['G_x'] for ori in orientations_data])
+        G_y = np.array([ori['G_y'] for ori in orientations_data])
+        G_z = np.array([ori['G_z'] for ori in orientations_data])
+        ids = np.array([ori['id'] for ori in orientations_data])
+        nugget = np.array([ori['nugget'] for ori in orientations_data])
+        
+        # Apply polarity to gradients
+        for i, ori in enumerate(orientations_data):
+            if ori['polarity'] == -1:
+                G_x[i] *= -1
+                G_y[i] *= -1
+                G_z[i] *= -1
+        
+        # Create name_id_map from unique IDs
+        unique_ids = np.unique(ids)
+        name_id_map = {f"surface_{id}": id for id in unique_ids}
+        
+        # Create OrientationsTable
+        return OrientationsTable.from_arrays(
+            x=x,
+            y=y,
+            z=z,
+            G_x=G_x,
+            G_y=G_y,
+            G_z=G_z,
+            names=[f"surface_{id}" for id in ids],
+            nugget=nugget,
+            name_id_map=name_id_map
+        )
 
     @staticmethod
     def save_model_to_json(model: GeoModel, file_path: str) -> None:
@@ -114,8 +178,8 @@ def _validate_json_schema(data: Dict[str, Any]) -> bool:
             bool: True if valid, False otherwise
         """
         # Check required top-level keys
-        required_keys = {'metadata', 'surface_points', 'orientations', 'faults', 
-                        'series', 'grid_settings', 'interpolation_options'}
+        required_keys = {'metadata', 'surface_points', 'orientations', 'series', 
+                        'grid_settings', 'interpolation_options'}
         if not all(key in data for key in required_keys):
             return False
 
@@ -132,4 +196,19 @@ def _validate_json_schema(data: Dict[str, Any]) -> bool:
             if not isinstance(sp['id'], int):
                 return False
 
+        # Validate orientations
+        if not isinstance(data['orientations'], list):
+            return False
+            
+        for ori in data['orientations']:
+            required_ori_keys = {'x', 'y', 'z', 'G_x', 'G_y', 'G_z', 'id', 'nugget', 'polarity'}
+            if not all(key in ori for key in required_ori_keys):
+                return False
+            if not all(isinstance(ori[key], (int, float)) for key in ['x', 'y', 'z', 'G_x', 'G_y', 'G_z', 'nugget']):
+                return False
+            if not isinstance(ori['id'], int):
+                return False
+            if not isinstance(ori['polarity'], int) or ori['polarity'] not in {-1, 1}:
+                return False
+                
         return True 

gempy/modules/json_io/schema.py

@@ -3,7 +3,7 @@
 This module defines the expected structure of JSON files for loading and saving GemPy models.
 """
 
-from typing import TypedDict, List, Dict, Any, Optional
+from typing import TypedDict, List, Dict, Any, Optional, Union, Sequence
 
 class SurfacePoint(TypedDict):
     x: float
@@ -16,23 +16,33 @@ class Orientation(TypedDict):
     x: float
     y: float
     z: float
-    G_x: float
-    G_y: float
-    G_z: float
+    G_x: float  # X component of the gradient
+    G_y: float  # Y component of the gradient
+    G_z: float  # Z component of the gradient
     id: int
-    polarity: int
+    nugget: float
+    polarity: int  # 1 for normal, -1 for reverse
+
+class Surface(TypedDict):
+    name: str
+    id: int
+    color: Optional[str]  # Hex color code
+    vertices: Optional[List[List[float]]]  # List of [x, y, z] coordinates
 
 class Fault(TypedDict):
     name: str
     id: int
     is_active: bool
+    surface: Surface
 
 class Series(TypedDict):
     name: str
     id: int
     is_active: bool
     is_fault: bool
     order_series: int
+    surfaces: List[Surface]
+    faults: List[Fault]
 
 class GridSettings(TypedDict):
     regular_grid_resolution: List[int]
@@ -49,7 +59,6 @@ class GemPyModelJson(TypedDict):
     metadata: ModelMetadata
     surface_points: List[SurfacePoint]
     orientations: List[Orientation]
-    faults: List[Fault]
     series: List[Series]
     grid_settings: GridSettings
     interpolation_options: Dict[str, Any] 

test/test_modules/test_json_io.py

@@ -35,9 +35,41 @@ def sample_surface_points():
 
 
 @pytest.fixture
-def sample_json_data(sample_surface_points):
+def sample_orientations():
+    """Create sample orientation data for testing."""
+    x = np.array([0.5, 1.5, 2.5, 3.5])
+    y = np.array([0.5, 1.5, 2.5, 3.5])
+    z = np.array([0.5, 1.5, 2.5, 3.5])
+    G_x = np.array([0, 0, 0, 0])
+    G_y = np.array([0, 0, 0, 0])
+    G_z = np.array([1, 1, -1, 1])  # One reversed orientation
+    ids = np.array([0, 1, 1, 2])  # Three different surfaces
+    nugget = np.array([0.01, 0.01, 0.01, 0.01])
+    
+    # Create name to id mapping
+    name_id_map = {f"surface_{id}": id for id in np.unique(ids)}
+    
+    # Create an OrientationsTable
+    orientations = gp.data.OrientationsTable.from_arrays(
+        x=x,
+        y=y,
+        z=z,
+        G_x=G_x,
+        G_y=G_y,
+        G_z=G_z,
+        names=[f"surface_{id}" for id in ids],
+        nugget=nugget,
+        name_id_map=name_id_map
+    )
+    
+    return orientations, x, y, z, G_x, G_y, G_z, ids, nugget, name_id_map
+
+
+@pytest.fixture
+def sample_json_data(sample_surface_points, sample_orientations):
     """Create sample JSON data for testing."""
-    _, x, y, z, ids, nugget, _ = sample_surface_points
+    _, x_sp, y_sp, z_sp, ids_sp, nugget_sp, _ = sample_surface_points
+    _, x_ori, y_ori, z_ori, G_x, G_y, G_z, ids_ori, nugget_ori, _ = sample_orientations
 
     return {
         "metadata": {
@@ -48,16 +80,28 @@ def sample_json_data(sample_surface_points):
         },
         "surface_points": [
             {
-                "x": float(x[i]),
-                "y": float(y[i]),
-                "z": float(z[i]),
-                "id": int(ids[i]),
-                "nugget": float(nugget[i])
+                "x": float(x_sp[i]),
+                "y": float(y_sp[i]),
+                "z": float(z_sp[i]),
+                "id": int(ids_sp[i]),
+                "nugget": float(nugget_sp[i])
             }
-            for i in range(len(x))
+            for i in range(len(x_sp))
+        ],
+        "orientations": [
+            {
+                "x": float(x_ori[i]),
+                "y": float(y_ori[i]),
+                "z": float(z_ori[i]),
+                "G_x": float(G_x[i]),
+                "G_y": float(G_y[i]),
+                "G_z": float(G_z[i]),
+                "id": int(ids_ori[i]),
+                "nugget": float(nugget_ori[i]),
+                "polarity": 1  # Always set to 1 since we're testing the raw G_z values
+            }
+            for i in range(len(x_ori))
         ],
-        "orientations": [],
-        "faults": [],
         "series": [],
         "grid_settings": {
             "regular_grid_resolution": [10, 10, 10],
@@ -84,6 +128,25 @@ def test_surface_points_loading(sample_surface_points, sample_json_data):
     assert surface_points.name_id_map == loaded_surface_points.name_id_map, "Name to ID mappings don't match"
 
 
+def test_orientations_loading(sample_orientations, sample_json_data):
+    """Test loading orientations from JSON data."""
+    orientations, _, _, _, _, _, _, _, _, name_id_map = sample_orientations
+    
+    # Load orientations from JSON
+    loaded_orientations = JsonIO._load_orientations(sample_json_data["orientations"])
+    
+    # Verify all data matches
+    assert np.allclose(orientations.xyz[:, 0], loaded_orientations.xyz[:, 0]), "X coordinates don't match"
+    assert np.allclose(orientations.xyz[:, 1], loaded_orientations.xyz[:, 1]), "Y coordinates don't match"
+    assert np.allclose(orientations.xyz[:, 2], loaded_orientations.xyz[:, 2]), "Z coordinates don't match"
+    assert np.allclose(orientations.grads[:, 0], loaded_orientations.grads[:, 0]), "G_x values don't match"
+    assert np.allclose(orientations.grads[:, 1], loaded_orientations.grads[:, 1]), "G_y values don't match"
+    assert np.allclose(orientations.grads[:, 2], loaded_orientations.grads[:, 2]), "G_z values don't match"
+    assert np.array_equal(orientations.ids, loaded_orientations.ids), "IDs don't match"
+    assert np.allclose(orientations.nugget, loaded_orientations.nugget), "Nugget values don't match"
+    assert orientations.name_id_map == loaded_orientations.name_id_map, "Name to ID mappings don't match"
+
+
 def test_surface_points_saving(tmp_path, sample_surface_points, sample_json_data):
     """Test saving surface points to JSON file."""
     surface_points, _, _, _, _, _, _ = sample_surface_points
@@ -131,4 +194,63 @@ def test_missing_surface_points_data():
     ]
 
     with pytest.raises(ValueError):
-        JsonIO._load_surface_points(invalid_data) 
+        JsonIO._load_surface_points(invalid_data)
+
+
+def test_invalid_orientations_data():
+    """Test handling of invalid orientation data."""
+    invalid_data = [
+        {
+            "x": 1.0,
+            "y": 1.0,
+            "z": 1.0,
+            "G_x": 0.0,
+            "G_y": 0.0,
+            "G_z": "invalid",  # Should be float
+            "id": 0,
+            "nugget": 0.01,
+            "polarity": 1
+        }
+    ]
+    
+    with pytest.raises(ValueError):
+        JsonIO._load_orientations(invalid_data)
+
+
+def test_missing_orientations_data():
+    """Test handling of missing orientation data."""
+    invalid_data = [
+        {
+            "x": 1.0,
+            "y": 1.0,
+            "z": 1.0,
+            "G_x": 0.0,
+            # Missing G_y and G_z
+            "id": 0,
+            "nugget": 0.01,
+            "polarity": 1
+        }
+    ]
+    
+    with pytest.raises(ValueError):
+        JsonIO._load_orientations(invalid_data)
+
+
+def test_invalid_orientation_polarity():
+    """Test handling of invalid orientation polarity."""
+    invalid_data = [
+        {
+            "x": 1.0,
+            "y": 1.0,
+            "z": 1.0,
+            "G_x": 0.0,
+            "G_y": 0.0,
+            "G_z": 1.0,
+            "id": 0,
+            "nugget": 0.01,
+            "polarity": 2  # Should be 1 or -1
+        }
+    ]
+    
+    with pytest.raises(ValueError):
+        JsonIO._load_orientations(invalid_data)