#!/usr/bin/python # Take a vertical slice in a 3D zyx array, along a varying y 'path' # # J-Y Peterschmitt - LSCE - 01/2015 import numpy as np # Specify the test variable shape z_nb_depth, y_nb_lat, x_nb_lon = 6, 5, 10 # Specify the values of the y indices for which we will take a # vertical slice. We need x_nb_lon values, from '0' to 'y_nb_lat-1' indices_for_y_slice = np.array([0, 1, 2, 3, 4, 3, 2, 1, 0, 1]) # Create a pseudo-variable with a content that we can easily # understand, and so that we can easily check if the slices we make # correct. All the axes have different sizes, so that we can easily # distinguish them # Create a vector along the y axis, and then replicate it along the # other axes to create the test variable y_vector = np.arange(3, 3+y_nb_lat) y_vector.shape = (y_nb_lat, 1) print('y_vector =') print(y_vector) zyx_variable = np.tile(y_vector, (z_nb_depth, 1, x_nb_lon)) print('\nzyx_variable shape =') print(zyx_variable.shape) print('\nzyx_variable =') print(zyx_variable) print('\nVertical zx slice for y_index = 1') print(zyx_variable[:, 1, :]) # For the vertical zx slice where we take a different value of y_index # for each value of x_index (y_index depends on x_index, NOT on # z_index), we need a way of specifying that we want a) ALL the values # along the x_lon axis, and b) ALL the values along the z_depth axis # # => we use the np.ix_ function to generate the indexing arrays! # # Note: for an example where y_index depends on x_index AND z_index, # see the indirect_indexing.py script zx_slice_tuple = np.ix_(range(z_nb_depth), range(x_nb_lon)) print('\nzx_slice_tuple = ') print(zx_slice_tuple) # We are ready to extract the slice! zx_advanced_slice = zyx_variable[zx_slice_tuple[0], indices_for_y_slice, zx_slice_tuple[1]] print('\nindices_for_y_slice =') print(indices_for_y_slice) print('\nzx_advanced_slice shape =') print(zx_advanced_slice.shape) print('\nzx_advanced_slice =') print(zx_advanced_slice) # Now we do the contrary, we assign a value to the slice! slicing_tuple = (zx_slice_tuple[0], indices_for_y_slice, zx_slice_tuple[1]) print('\n*** Assigning a constant value to the slice... ***') zyx_variable[slicing_tuple] = -1 print('\nNew value of zyx_variable[slicing_tuple] =') print(zyx_variable[slicing_tuple]) print('\nNew value of zyx_variable = ') print(zyx_variable) # The end