33import numpy as np
44import pytest
55
6- from mdio .converters .segy import _calculate_live_mask_chunksize
7- from mdio .core import Grid , Dimension
86from mdio .constants import INT32_MAX
7+ from mdio .converters .segy import _calculate_live_mask_chunksize
8+ from mdio .core import Dimension
9+ from mdio .core import Grid
910
1011
1112def test_small_grid_no_chunking ():
@@ -14,11 +15,11 @@ def test_small_grid_no_chunking():
1415 dims = [
1516 Dimension (coords = range (0 , 100 , 1 ), name = "dim1" ),
1617 Dimension (coords = range (0 , 100 , 1 ), name = "dim2" ),
17- Dimension (coords = range (0 , 100 , 1 ), name = "sample" )
18+ Dimension (coords = range (0 , 100 , 1 ), name = "sample" ),
1819 ]
1920 grid = Grid (dims = dims )
2021 grid .live_mask = np .ones ((100 , 100 ), dtype = bool )
21-
22+
2223 result = _calculate_live_mask_chunksize (grid )
2324 assert result == - 1
2425
@@ -30,19 +31,19 @@ def test_large_2d_grid_chunking():
3031 dims = [
3132 Dimension (coords = range (0 , 50000 , 1 ), name = "dim1" ),
3233 Dimension (coords = range (0 , 50000 , 1 ), name = "dim2" ),
33- Dimension (coords = range (0 , 100 , 1 ), name = "sample" )
34+ Dimension (coords = range (0 , 100 , 1 ), name = "sample" ),
3435 ]
3536 grid = Grid (dims = dims )
3637 grid .live_mask = np .ones ((50000 , 50000 ), dtype = bool )
37-
38+
3839 result = _calculate_live_mask_chunksize (grid )
39-
40+
4041 # Calculate expected values
4142 total_elements = 50000 * 50000
4243 num_chunks = np .ceil (total_elements / INT32_MAX ).astype (int )
43- dim_chunks = int (np .ceil (50000 / np .ceil (np .power (num_chunks , 1 / 2 ))))
44+ dim_chunks = int (np .ceil (50000 / np .ceil (np .power (num_chunks , 1 / 2 ))))
4445 expected_chunk_size = int (np .ceil (50000 / dim_chunks ))
45-
46+
4647 assert result == (expected_chunk_size , expected_chunk_size )
4748
4849
@@ -54,19 +55,19 @@ def test_large_3d_grid_chunking():
5455 Dimension (coords = range (0 , 1500 , 1 ), name = "dim1" ),
5556 Dimension (coords = range (0 , 1500 , 1 ), name = "dim2" ),
5657 Dimension (coords = range (0 , 1500 , 1 ), name = "dim3" ),
57- Dimension (coords = range (0 , 100 , 1 ), name = "sample" )
58+ Dimension (coords = range (0 , 100 , 1 ), name = "sample" ),
5859 ]
5960 grid = Grid (dims = dims )
6061 grid .live_mask = np .ones ((1500 , 1500 , 1500 ), dtype = bool )
61-
62+
6263 result = _calculate_live_mask_chunksize (grid )
63-
64+
6465 # Calculate expected values
6566 total_elements = 1500 * 1500 * 1500
6667 num_chunks = np .ceil (total_elements / INT32_MAX ).astype (int )
67- dim_chunks = int (np .ceil (1500 / np .ceil (np .power (num_chunks , 1 / 3 ))))
68+ dim_chunks = int (np .ceil (1500 / np .ceil (np .power (num_chunks , 1 / 3 ))))
6869 expected_chunk_size = int (np .ceil (1500 / dim_chunks ))
69-
70+
7071 assert result == (expected_chunk_size , expected_chunk_size , expected_chunk_size )
7172
7273
@@ -78,19 +79,19 @@ def test_uneven_dimensions_chunking():
7879 dims = [
7980 Dimension (coords = range (0 , 50000 , 1 ), name = "dim1" ),
8081 Dimension (coords = range (0 , 50000 , 1 ), name = "dim2" ),
81- Dimension (coords = range (0 , 100 , 1 ), name = "sample" )
82+ Dimension (coords = range (0 , 100 , 1 ), name = "sample" ),
8283 ]
8384 grid = Grid (dims = dims )
8485 grid .live_mask = np .ones ((50000 , 50000 ), dtype = bool )
85-
86+
8687 result = _calculate_live_mask_chunksize (grid )
87-
88+
8889 # Calculate expected values
8990 total_elements = 50000 * 50000
9091 num_chunks = np .ceil (total_elements / INT32_MAX ).astype (int )
91- dim_chunks = int (np .ceil (50000 / np .ceil (np .power (num_chunks , 1 / 2 ))))
92+ dim_chunks = int (np .ceil (50000 / np .ceil (np .power (num_chunks , 1 / 2 ))))
9293 expected_chunk_size = int (np .ceil (50000 / dim_chunks ))
93-
94+
9495 assert result == (expected_chunk_size , expected_chunk_size )
9596
9697
@@ -108,33 +109,38 @@ def test_prestack_land_survey_chunking():
108109 Dimension (coords = range (0 , 1000 , 1 ), name = "receiver_point" ),
109110 Dimension (coords = range (0 , 100 , 1 ), name = "offset" ),
110111 Dimension (coords = range (0 , 36 , 1 ), name = "azimuth" ),
111- Dimension (coords = range (0 , 1000 , 1 ), name = "sample" )
112+ Dimension (coords = range (0 , 1000 , 1 ), name = "sample" ),
112113 ]
113114 grid = Grid (dims = dims )
114115 grid .live_mask = np .ones ((1000 , 1000 , 100 , 36 ), dtype = bool )
115-
116+
116117 result = _calculate_live_mask_chunksize (grid )
117-
118+
118119 # Calculate expected values
119120 total_elements = 1000 * 1000 * 100 * 36
120121 num_chunks = np .ceil (total_elements / INT32_MAX ).astype (int )
121- dim_chunks = int (np .ceil (1000 / np .ceil (np .power (num_chunks , 1 / 4 ))))
122+ dim_chunks = int (np .ceil (1000 / np .ceil (np .power (num_chunks , 1 / 4 ))))
122123 expected_chunk_size = int (np .ceil (1000 / dim_chunks ))
123-
124+
124125 # For a 4D grid, we expect chunk sizes to be distributed across all dimensions
125126 # The chunk size should be the same for all dimensions since they're all equally important
126- assert result == (expected_chunk_size , expected_chunk_size , expected_chunk_size , expected_chunk_size )
127+ assert result == (
128+ expected_chunk_size ,
129+ expected_chunk_size ,
130+ expected_chunk_size ,
131+ expected_chunk_size ,
132+ )
127133
128134
129135def test_edge_case_empty_grid ():
130136 """Test empty grid edge case."""
131137 dims = [
132138 Dimension (coords = range (0 , 0 , 1 ), name = "dim1" ),
133139 Dimension (coords = range (0 , 0 , 1 ), name = "dim2" ),
134- Dimension (coords = range (0 , 100 , 1 ), name = "sample" )
140+ Dimension (coords = range (0 , 100 , 1 ), name = "sample" ),
135141 ]
136142 grid = Grid (dims = dims )
137143 grid .live_mask = np .zeros ((0 , 0 ), dtype = bool )
138-
144+
139145 result = _calculate_live_mask_chunksize (grid )
140- assert result == - 1 # Empty grid shouldn't need chunking
146+ assert result == - 1 # Empty grid shouldn't need chunking
0 commit comments