chore: reverts to_na changes and refactors write_xy_narray to use Num…

…o::NArray instead
syngenta · Nov 27, 2024 · 39bea8c · 39bea8c
1 parent c87e87f
commit 39bea8c
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 4 deletions.
diff --git a/lib/gdal/extensions/raster_band/extensions.rb b/lib/gdal/extensions/raster_band/extensions.rb
@@ -30,7 +30,7 @@ def to_a
       #
       # @return [NArray]
       def to_na(to_data_type = nil)
-        narray = to_nna
+        narray = NArray.to_na(to_a)
 
         return narray unless to_data_type
 
@@ -39,8 +39,11 @@ def to_na(to_data_type = nil)
         narray.to_type(narray_type)
       end
 
+      # Iterates through all lines and builds an NArray of pixels.
+      #
+      # @return [Numo::NArray]
       def to_nna
-        NArray.to_na(to_a)
+        Numo::NArray[*to_a]
       end
 
       # Each pixel of the raster projected using the dataset's geo_transform.

diff --git a/lib/gdal/extensions/raster_band/io_extensions.rb b/lib/gdal/extensions/raster_band/io_extensions.rb
@@ -26,11 +26,14 @@ def readlines
       # @param pixel_array [NArray] The 2d list of pixels.
       def write_xy_narray(pixel_array)
         data_pointer = FFI::MemoryPointer.new(:buffer_out, block_buffer_size)
+        nnarray = GDAL._gdal_data_type_to_numo_narray_type_constant(data_type)
+
+        pixel_array = nnarray[*pixel_array.to_a]
 
         block_count[:y].times do |y_block_number|
           block_count[:x].times do |x_block_number|
             # Create a block-sized NArray of zeros. This represents the block of pixels that will be written.
-            block_pixels = NArray.to_na(Array.new(block_size[:y]) { Array.new(block_size[:x], 0) })
+            block_pixels = nnarray.zeros(block_size[:y], block_size[:x])
 
             y_block_size = calculate_y_block_size(y_block_number)
             x_block_size = calculate_x_block_size(x_block_number)
@@ -40,7 +43,7 @@ def write_xy_narray(pixel_array)
             source_y_range = y_block_number * block_size[:y]...(y_block_number * block_size[:y]) + y_block_size
 
             # Copy the corresponding pixels from the input array to the block pixels.
-            block_pixels[0...x_block_size, 0...y_block_size] = pixel_array[source_x_range, source_y_range]
+            block_pixels[0...y_block_size, 0...x_block_size] = pixel_array[source_y_range, source_x_range]
 
             GDAL._write_pointer(data_pointer, data_type, block_pixels.to_a.flatten)
 

diff --git a/lib/gdal/internal_helpers.rb b/lib/gdal/internal_helpers.rb
@@ -193,6 +193,23 @@ def _gdal_data_type_to_narray_type_constant(data_type)
         end
       end
 
+      # Maps GDAL DataTypes to [Numo::NArray] type constants.
+      #
+      # @param data_type [FFI::GDAL::GDAL::DataType]
+      # @return [Class]
+      def _gdal_data_type_to_numo_narray_type_constant(data_type)
+        case data_type
+        when :GDT_Byte                                              then Numo::UInt8
+        when :GDT_Int16                                             then Numo::Int16
+        when :GDT_UInt16, :GDT_Int32, :GDT_UInt32                   then Numo::Int32
+        when :GDT_Float32                                           then Numo::SFloat
+        when :GDT_Float64                                           then Numo::DFloat
+        when :GDT_CInt16, :GDT_CInt32, :GDT_CFloat32, :GDT_CFloat64 then Numo::SComplex
+        else
+          raise GDAL::InvalidDataType, "Unknown data type: #{data_type}"
+        end
+      end
+
       # Helper method for reading an FFI pointer based on the GDAL DataType of
       # the pointer.
       #