def test_rf_local_is_in(self):
from pyspark.sql.functions import lit, array, col
from pyspark.sql import Row
nd = 5
t = Tile(np.array([[1, 3, 4], [nd, 0, 3]]),
CellType.uint8().with_no_data_value(nd))
# note the convert is due to issue #188
df = self.spark.createDataFrame([Row(t=t)]) \
.withColumn('a', array(lit(3), lit(4))) \
.withColumn('in2', rf_convert_cell_type(
rf_local_is_in(col('t'), array(lit(0), lit(4))),
'uint8')) \
.withColumn('in3', rf_convert_cell_type(rf_local_is_in('t', 'a'), 'uint8')) \
.withColumn('in4', rf_convert_cell_type(
rf_local_is_in('t', array(lit(0), lit(4), lit(3))),
'uint8')) \
.withColumn('in_list', rf_convert_cell_type(rf_local_is_in(col('t'), [4, 1]), 'uint8'))
result = df.first()
self.assertEqual(result['in2'].cells.sum(), 2)
assert_equal(result['in2'].cells, np.isin(t.cells, np.array([0, 4])))
self.assertEqual(result['in3'].cells.sum(), 3)
self.assertEqual(result['in4'].cells.sum(), 4)
self.assertEqual(
result['in_list'].cells.sum(), 2,
"Tile value {} should contain two 1s as: [[1, 0, 1],[0, 0, 0]]".
format(result['in_list'].cells))
def test_mask(self):
from pyspark.sql import Row
from pyrasterframes.rf_types import Tile, CellType
np.random.seed(999)
# importantly exclude 0 from teh range because that's the nodata value for the `data_tile`'s cell type
ma = np.ma.array(np.random.randint(1, 10, (5, 5), dtype='int8'),
mask=np.random.rand(5, 5) > 0.7)
expected_data_values = ma.compressed().size
expected_no_data_values = ma.size - expected_data_values
self.assertTrue(expected_data_values > 0,
"Make sure random seed is cooperative ")
self.assertTrue(expected_no_data_values > 0,
"Make sure random seed is cooperative ")
data_tile = Tile(np.ones(ma.shape, ma.dtype), CellType.uint8())
df = self.spark.createDataFrame([Row(t=data_tile, m=Tile(ma))]) \
.withColumn('masked_t', rf_mask('t', 'm'))
result = df.select(rf_data_cells('masked_t')).first()[0]
self.assertEqual(
result, expected_data_values,
f"Masked tile should have {expected_data_values} data values but found: {df.select('masked_t').first()[0].cells}."
f"Original data: {data_tile.cells}"
f"Masked by {ma}")
nd_result = df.select(rf_no_data_cells('masked_t')).first()[0]
self.assertEqual(nd_result, expected_no_data_values)
# deser of tile is correct
self.assertEqual(
df.select('masked_t').first()[0].cells.compressed().size,
expected_data_values)
def test_rf_where(self):
cond = Tile(np.random.binomial(1, 0.35, (10, 10)), CellType.uint8())
x = Tile(np.random.randint(-20, 10, (10, 10)), CellType.int8())
y = Tile(np.random.randint(0, 30, (10, 10)), CellType.int8())
df = self.spark.createDataFrame([Row(cond=cond, x=x, y=y)])
result = df.select(rf_where('cond', 'x', 'y')).first()[0].cells
assert_equal(result, np.where(cond.cells, x.cells, y.cells))
def test_mask_by_values(self):
tile = Tile(np.random.randint(1, 100, (5, 5)), CellType.uint8())
mask_tile = Tile(np.array(range(1, 26), 'uint8').reshape(5, 5))
expected_diag_nd = Tile(np.ma.masked_array(tile.cells, mask=np.eye(5)))
df = self.spark.createDataFrame([Row(t=tile, m=mask_tile)]) \
.select(rf_mask_by_values('t', 'm', [0, 6, 12, 18, 24])) # values on the diagonal
result0 = df.first()
# assert_equal(result0[0].cells, expected_diag_nd)
self.assertTrue(result0[0] == expected_diag_nd)
def test_rf_local_data_and_no_data(self):
from pyspark.sql import Row
from pyrasterframes.rf_types import Tile
nd = 5
t = Tile(np.array([[1, 3, 4], [nd, 0, 3]]),
CellType.uint8().with_no_data_value(nd))
# note the convert is due to issue #188
df = self.spark.createDataFrame([Row(t=t)])\
.withColumn('lnd', rf_convert_cell_type(rf_local_no_data('t'), 'uint8')) \
.withColumn('ld', rf_convert_cell_type(rf_local_data('t'), 'uint8'))
result = df.first()
result_nd = result['lnd']
assert_equal(result_nd.cells, t.cells.mask)
result_d = result['ld']
assert_equal(result_d.cells, np.invert(t.cells.mask))
def test_rf_interpret_cell_type_as(self):
from pyspark.sql import Row
from pyrasterframes.rf_types import Tile
df = self.spark.createDataFrame([
Row(t=Tile(np.array([[1, 3, 4], [5, 0, 3]]),
CellType.uint8().with_no_data_value(5)))
])
df = df.withColumn('tile', rf_interpret_cell_type_as(
't', 'uint8ud3')) # threes become ND
result = df.select(
rf_tile_sum(rf_local_equal(
't', lit(3))).alias('threes')).first()['threes']
self.assertEqual(result, 2)
result_5 = df.select(
rf_tile_sum(rf_local_equal(
't', lit(5))).alias('fives')).first()['fives']
self.assertEqual(result_5, 0)
