fixed issues with to_array function

Author: prasad-sawantdesai

imas/ids_slice.py

@@ -77,10 +77,10 @@ def _path(self) -> str:
     @property
     def shape(self) -> Tuple[int, ...]:
         """Get the virtual multi-dimensional shape.
-        
+
         Returns the shape of the data as if it were organized in a multi-dimensional
         array, based on the hierarchy of slicing operations performed.
-        
+
         Returns:
             Tuple of dimensions. Use with caution for jagged arrays where sizes vary.
         """
@@ -119,7 +119,7 @@ def __getitem__(self, item: Union[int, slice]) -> Union[Any, "IDSSlice"]:
                 # Preserve structure instead of flattening
                 sliced_elements = []
                 sliced_sizes = []
-                
+
                 for array in self._matched_elements:
                     sliced = array[item]
                     if isinstance(sliced, IDSSlice):
@@ -134,7 +134,9 @@ def __getitem__(self, item: Union[int, slice]) -> Union[Any, "IDSSlice"]:
 
                 # Update shape to reflect the sliced structure
                 # Keep first dimensions, update last dimension
-                new_virtual_shape = self._virtual_shape[:-1] + (sliced_sizes[0] if sliced_sizes else 0,)
+                new_virtual_shape = self._virtual_shape[:-1] + (
+                    sliced_sizes[0] if sliced_sizes else 0,
+                )
                 new_hierarchy = self._element_hierarchy[:-1] + [sliced_sizes]
 
                 return IDSSlice(
@@ -172,7 +174,9 @@ def __getitem__(self, item: Union[int, slice]) -> Union[Any, "IDSSlice"]:
 
                 # Update shape to reflect the slice on first dimension
                 new_virtual_shape = (len(sliced_elements),) + self._virtual_shape[1:]
-                new_element_hierarchy = [len(sliced_elements)] + self._element_hierarchy[1:]
+                new_element_hierarchy = [
+                    len(sliced_elements)
+                ] + self._element_hierarchy[1:]
 
                 return IDSSlice(
                     self.metadata,
@@ -232,11 +236,13 @@ def __getattr__(self, name: str) -> "IDSSlice":
         if isinstance(child_elements[0], IDSStructArray):
             # Children are IDSStructArray - track the new dimension
             child_sizes = [len(arr) for arr in child_elements]
-            
+
             # New virtual shape: current shape + new dimension
-            new_virtual_shape = self._virtual_shape + (child_sizes[0] if child_sizes else 0,)
+            new_virtual_shape = self._virtual_shape + (
+                child_sizes[0] if child_sizes else 0,
+            )
             new_hierarchy = self._element_hierarchy + [child_sizes]
-            
+
             return IDSSlice(
                 child_metadata,
                 child_elements,
@@ -249,12 +255,14 @@ def __getattr__(self, name: str) -> "IDSSlice":
             # Children are IDSNumericArray - track the array dimension
             # Each IDSNumericArray has a size (length of its data)
             child_sizes = [len(arr) for arr in child_elements]
-            
-            # New virtual shape: current shape + new dimension (size of first numeric array)
+
+            # New virtual shape: current shape + new dimension
             # Jagged arrays handled by to_array() with object dtype
-            new_virtual_shape = self._virtual_shape + (child_sizes[0] if child_sizes else 0,)
+            new_virtual_shape = self._virtual_shape + (
+                child_sizes[0] if child_sizes else 0,
+            )
             new_hierarchy = self._element_hierarchy + [child_sizes]
-            
+
             return IDSSlice(
                 child_metadata,
                 child_elements,
@@ -338,7 +346,11 @@ def values(self, reshape: bool = False) -> Any:
             >>> # Result: ndarray shape (106, 100)
             >>>
             >>> # 3D ions case - returns object array with structure
-            >>> ion_rho = core_profiles.profiles_1d[:].ion[:].element[:].density.values(reshape=True)
+            >>> ion_rho = (
+            ...     core_profiles.profiles_1d[:].ion[:].element[:].density.values(
+            ...         reshape=True
+            ...     )
+            ... )
             >>> # Result: object array shape (106, 3, 2) with IDSNumericArray elements
         """
         from imas.ids_primitive import IDSPrimitive, IDSNumericArray
@@ -359,7 +371,9 @@ def values(self, reshape: bool = False) -> Any:
             if isinstance(element, IDSPrimitive):
                 flat_values.append(element.value)
             elif isinstance(element, IDSNumericArray):
-                flat_values.append(element.data if hasattr(element, 'data') else element.value)
+                flat_values.append(
+                    element.data if hasattr(element, "data") else element.value
+                )
             else:
                 flat_values.append(element)
 
@@ -399,7 +413,8 @@ def to_array(self) -> np.ndarray:
         structure of the IMAS data.
 
         Returns:
-            numpy.ndarray with shape self.shape. For jagged arrays, dtype will be object.
+            numpy.ndarray with shape self.shape. For jagged arrays,
+            dtype will be object.
 
         Raises:
             ValueError: If array cannot be converted to numpy
@@ -436,25 +451,27 @@ def to_array(self) -> np.ndarray:
 
         # Multi-dimensional case
         # Check if matched elements are themselves arrays (IDSNumericArray)
-        if self._matched_elements and isinstance(self._matched_elements[0], IDSNumericArray):
+        if self._matched_elements and isinstance(
+            self._matched_elements[0], IDSNumericArray
+        ):
             # Elements are numeric arrays - extract their values and stack them
             array_values = []
             for element in self._matched_elements:
                 if isinstance(element, IDSNumericArray):
                     array_values.append(element.value)
                 else:
                     array_values.append(element)
-            
+
             # Try to stack into proper shape
             try:
                 # Check if all arrays have the same size (regular)
                 sizes = []
                 for val in array_values:
-                    if hasattr(val, '__len__'):
+                    if hasattr(val, "__len__"):
                         sizes.append(len(val))
                     else:
                         sizes.append(1)
-                
+
                 # If all sizes are the same, we can create a regular array
                 if len(set(sizes)) == 1:
                     # Regular array - all sub-arrays same size
@@ -478,7 +495,7 @@ def to_array(self) -> np.ndarray:
                     for i, val in enumerate(array_values):
                         result_arr[i] = val
                     return result_arr
-            except (ValueError, TypeError) as e:
+            except (ValueError, TypeError):
                 # Fallback: return object array
                 result_arr = np.empty(self._virtual_shape[0], dtype=object)
                 for i, val in enumerate(array_values):
@@ -488,11 +505,29 @@ def to_array(self) -> np.ndarray:
         # For non-numeric elements in multi-dimensional structure
         # Extract and try to build structure
         flat_values = []
-        for element in self._matched_elements:
-            if isinstance(element, IDSPrimitive):
-                flat_values.append(element.value)
-            else:
-                flat_values.append(element)
+
+        # First check if matched_elements are IDSStructArray (which need flattening)
+        from imas.ids_struct_array import IDSStructArray
+
+        has_struct_arrays = self._matched_elements and isinstance(
+            self._matched_elements[0], IDSStructArray
+        )
+
+        if has_struct_arrays:
+            # Flatten IDSStructArray elements
+            for struct_array in self._matched_elements:
+                for element in struct_array:
+                    if isinstance(element, IDSPrimitive):
+                        flat_values.append(element.value)
+                    else:
+                        flat_values.append(element)
+        else:
+            # Regular elements
+            for element in self._matched_elements:
+                if isinstance(element, IDSPrimitive):
+                    flat_values.append(element.value)
+                else:
+                    flat_values.append(element)
 
         total_size = 1
         for dim in self._virtual_shape:

imas/test/test_multidim_slicing.py

@@ -6,7 +6,6 @@
 import pytest
 
 from imas.ids_factory import IDSFactory
-from imas.ids_slice import IDSSlice
 
 
 class TestMultiDimSlicing:
@@ -137,7 +136,7 @@ def test_slice_preserves_groups(self):
 
         # Get all ions, then slice
         result = cp.profiles_1d[:].ion[:]
-        
+
         # Should still know the structure: 10 profiles, 3 ions each
         assert result.shape == (10, 3)
         assert len(result) == 30  # Flattened for iteration, but shape preserved
@@ -153,7 +152,7 @@ def test_integer_index_on_nested(self):
 
         # Get first ion from all profiles
         result = cp.profiles_1d[:].ion[0]
-        
+
         assert len(result) == 5
         for i, ion in enumerate(result):
             assert ion.label == f"ion_{i}_0"
@@ -167,7 +166,7 @@ def test_slice_on_nested_arrays(self):
 
         # Get first 2 ions from each profile
         result = cp.profiles_1d[:].ion[:2]
-        
+
         assert result.shape == (5, 2)
         assert len(result) == 10  # 5 profiles * 2 ions each
 
@@ -180,7 +179,7 @@ def test_step_slicing_on_nested(self):
 
         # Get every other ion
         result = cp.profiles_1d[:].ion[::2]
-        
+
         assert result.shape == (5, 3)  # 5 profiles, 3 ions each (0, 2, 4)
         assert len(result) == 15
 
@@ -195,7 +194,7 @@ def test_negative_indexing_on_nested(self):
 
         # Get last ion from each profile
         result = cp.profiles_1d[:].ion[-1]
-        
+
         assert len(result) == 5
         for ion in result:
             assert ion.label == "ion_2"
@@ -227,7 +226,7 @@ def test_boolean_indexing_simple(self):
             p.electrons.density = np.array([float(i)] * 5)
 
         result = cp.profiles_1d[:].electrons.density
-        
+
         mask = np.array([True, False, True, False, True])
         filtered = result[mask]
         assert len(filtered) == 3
@@ -239,15 +238,10 @@ def test_assignment_on_slice(self):
         for p in cp.profiles_1d:
             p.grid.rho_tor_norm = np.array([0.0, 0.5, 1.0])
 
-        # Assign new values through slice
-        new_values = np.array([[0.1, 0.6, 1.1], 
-                               [0.2, 0.7, 1.2],
-                               [0.3, 0.8, 1.3]])
-        
         # This requires assignment support
         # cp.profiles_1d[:].grid.rho_tor_norm[:] = new_values
         # For now, verify slicing works for reading
-        
+
         result = cp.profiles_1d[:].grid.rho_tor_norm
         array = result.to_array()
         assert array.shape == (3, 3)
@@ -257,7 +251,7 @@ def test_xarray_integration_compatible(self):
         cp = IDSFactory("3.39.0").core_profiles()
         cp.profiles_1d.resize(3)
         cp.time = np.array([1.0, 2.0, 3.0])
-        
+
         for i, p in enumerate(cp.profiles_1d):
             p.grid.rho_tor_norm = np.array([0.0, 0.5, 1.0])
             p.electrons.temperature = np.array([1.0, 2.0, 3.0]) * (i + 1)
@@ -274,7 +268,7 @@ def test_performance_large_hierarchy(self):
         cp = IDSFactory("3.39.0").core_profiles()
         n_profiles = 50
         cp.profiles_1d.resize(n_profiles)
-        
+
         for p in cp.profiles_1d:
             p.grid.rho_tor_norm = np.linspace(0, 1, 100)
             p.ion.resize(5)
@@ -284,7 +278,7 @@ def test_performance_large_hierarchy(self):
         # Should handle large data without significant slowdown
         result = cp.profiles_1d[:].grid.rho_tor_norm
         array = result.to_array()
-        
+
         assert array.shape == (n_profiles, 100)
 
     def test_lazy_loading_with_multidim(self):
@@ -296,12 +290,10 @@ def test_lazy_loading_with_multidim(self):
             p.grid.rho_tor_norm = np.array([0.0, 0.5, 1.0])
 
         result = cp.profiles_1d[:].grid.rho_tor_norm
-        
-        # Verify lazy attributes are preserved
-        assert hasattr(result, '_lazy')
-        assert hasattr(result, '_parent_array')
-
 
+        # Verify lazy attributes are preserved
+        assert hasattr(result, "_lazy")
+        assert hasattr(result, "_parent_array")
 
 
 class TestEdgeCases:
@@ -339,7 +331,7 @@ def test_single_dimension_value(self):
                 i.z_ion = 1.0
 
         result = cp.profiles_1d[:].ion[0].z_ion
-        
+
         # Should be 3 items (one per profile)
         assert len(result) == 3
 
@@ -352,6 +344,6 @@ def test_slice_of_slice(self):
 
         result1 = cp.profiles_1d[::2].ion  # Every other profile's ions
         assert result1.shape == (5, 3)
-        
+
         result2 = result1[:2]  # First 2 from each
         assert result2.shape == (5, 2)