Example #1
0
def _plot_start_or_end_markers(
        storm_object_table_one_track, axes_object, plot_at_start, marker_type,
        marker_size, track_colour=None, colour_map_object=None,
        colour_norm_object=None):
    """Plots marker at beginning or end of each storm track.

    :param storm_object_table_one_track: Same as input for `plot_storm_tracks`,
        except that this table contains only one track (primary storm ID).
    :param axes_object: See doc for `plot_storm_outlines`.
    :param plot_at_start: Boolean flag.  If True, will plot at beginning of each
        track.  If False, at end of each track.
    :param marker_type: Marker type.
    :param marker_size: Marker size.
    :param track_colour: Track colour.  This may be None.
    :param colour_map_object: [used only if `track_colour is None`]:
        Colour scheme (instance of `matplotlib.pyplot.cm` or similar).
    :param colour_norm_object: Normalizer for colour scheme.
    """

    num_storm_objects = len(storm_object_table_one_track.index)

    for i in range(num_storm_objects):
        if plot_at_start:
            these_indices = temporal_tracking.find_immediate_predecessors(
                storm_object_table=storm_object_table_one_track, target_row=i
            )
        else:
            these_indices = temporal_tracking.find_immediate_successors(
                storm_object_table=storm_object_table_one_track, target_row=i
            )

        if len(these_indices) > 0:
            continue

        if colour_map_object is None:
            this_colour = track_colour
        else:
            this_time_unix_sec = storm_object_table_one_track[
                tracking_utils.VALID_TIME_COLUMN
            ].values[i]

            this_colour = colour_map_object(colour_norm_object(
                this_time_unix_sec
            ))

        this_x_coord = storm_object_table_one_track[
            tracking_utils.CENTROID_X_COLUMN
        ].values[i]

        this_y_coord = storm_object_table_one_track[
            tracking_utils.CENTROID_Y_COLUMN
        ].values[i]

        axes_object.plot(
            this_x_coord, this_y_coord, linestyle='None',
            marker=marker_type, markersize=marker_size,
            markerfacecolor=this_colour, markeredgecolor=this_colour,
            markeredgewidth=2 if marker_type == 'x' else 1
        )
    def test_find_immediate_predecessors_first_5sec(self):
        """Ensures correct output from find_immediate_predecessors.

        In this case, working on the first table with max time difference =
        5 seconds.
        """

        this_num_storm_objects = len(FIRST_STORM_OBJECT_TABLE.index)

        for i in range(this_num_storm_objects):
            these_predecessor_rows = (
                temporal_tracking.find_immediate_predecessors(
                    storm_object_table=FIRST_STORM_OBJECT_TABLE,
                    target_row=i,
                    max_time_diff_seconds=5))

            these_predecessor_rows = numpy.sort(these_predecessor_rows)
            these_expected_rows = numpy.sort(
                numpy.array(FIRST_IMMED_PREDECESSOR_DICT_5SEC[i], dtype=int))

            self.assertTrue(
                numpy.array_equal(these_predecessor_rows, these_expected_rows))
Example #3
0
def plot_storm_tracks(storm_object_table,
                      axes_object,
                      basemap_object,
                      colour_map_object='random',
                      line_colour=DEFAULT_TRACK_COLOUR,
                      line_width=DEFAULT_TRACK_WIDTH,
                      start_marker_type=DEFAULT_START_MARKER_TYPE,
                      end_marker_type=DEFAULT_END_MARKER_TYPE,
                      start_marker_size=DEFAULT_START_MARKER_SIZE,
                      end_marker_size=DEFAULT_END_MARKER_SIZE):
    """Plots one or more storm tracks on the same map.

    :param storm_object_table: See doc for `plot_storm_outlines`.
    :param axes_object: Same.
    :param basemap_object: Same.
    :param colour_map_object: There are 3 cases.

    If "random", each track will be plotted in a random colour from
    `get_storm_track_colours`.

    If None, each track will be plotted in `line_colour` (the next input arg).

    If real colour map (instance of `matplotlib.pyplot.cm`), track segments will
    be coloured by time, according to this colour map.

    :param line_colour: [used only if `colour_map_object is None`]
        length-3 numpy array with (R, G, B).  Will be used for all tracks.
    :param line_width: Width of each storm track.
    :param start_marker_type: Marker type for beginning of track (in any format
        accepted by `matplotlib.lines`).  If `start_marker_type is None`,
        markers will not be used to show beginning of each track.
    :param end_marker_type: Same but for end of track.
    :param start_marker_size: Size of each start-point marker.
    :param end_marker_size: Size of each end-point marker.
    """

    plot_start_markers = start_marker_type is not None
    plot_end_markers = end_marker_type is not None

    if start_marker_type is None:
        start_marker_type = DEFAULT_START_MARKER_TYPE
        start_marker_size = DEFAULT_START_MARKER_SIZE

    if end_marker_type is None:
        end_marker_type = DEFAULT_END_MARKER_TYPE
        end_marker_size = DEFAULT_END_MARKER_SIZE

    x_coords_metres, y_coords_metres = basemap_object(
        storm_object_table[tracking_utils.CENTROID_LONGITUDE_COLUMN].values,
        storm_object_table[tracking_utils.CENTROID_LATITUDE_COLUMN].values)

    storm_object_table = storm_object_table.assign(
        **{
            tracking_utils.CENTROID_X_COLUMN: x_coords_metres,
            tracking_utils.CENTROID_Y_COLUMN: y_coords_metres
        })

    rgb_matrix = None
    num_colours = None
    colour_norm_object = None

    if colour_map_object is None:
        error_checking.assert_is_numpy_array(line_colour,
                                             exact_dimensions=numpy.array(
                                                 [3], dtype=int))

        rgb_matrix = numpy.reshape(line_colour, (1, 3))
        num_colours = rgb_matrix.shape[0]

    elif colour_map_object == 'random':
        rgb_matrix = get_storm_track_colours()
        num_colours = rgb_matrix.shape[0]

        colour_map_object = None

    else:
        first_time_unix_sec = numpy.min(
            storm_object_table[tracking_utils.VALID_TIME_COLUMN].values)
        last_time_unix_sec = numpy.max(
            storm_object_table[tracking_utils.VALID_TIME_COLUMN].values)

        colour_norm_object = pyplot.Normalize(first_time_unix_sec,
                                              last_time_unix_sec)

    track_primary_id_strings, object_to_track_indices = numpy.unique(
        storm_object_table[tracking_utils.PRIMARY_ID_COLUMN].values,
        return_inverse=True)

    num_tracks = len(track_primary_id_strings)

    for k in range(num_tracks):
        if colour_map_object is None:
            this_colour = rgb_matrix[numpy.mod(k, num_colours), :]
            this_colour = plotting_utils.colour_from_numpy_to_tuple(
                this_colour)
        else:
            this_colour = None

        these_object_indices = numpy.where(object_to_track_indices == k)[0]

        for i in these_object_indices:
            these_next_indices = temporal_tracking.find_immediate_successors(
                storm_object_table=storm_object_table, target_row=i)

            # if len(these_next_indices) > 1:
            #     axes_object.text(
            #         storm_object_table[
            #             tracking_utils.CENTROID_X_COLUMN].values[i],
            #         storm_object_table[
            #             tracking_utils.CENTROID_Y_COLUMN].values[i],
            #         '{0:d}-WAY SPLIT'.format(len(these_next_indices)),
            #         fontsize=12, color='k',
            #         horizontalalignment='left', verticalalignment='top')

            for j in these_next_indices:
                these_x_coords_metres = storm_object_table[
                    tracking_utils.CENTROID_X_COLUMN].values[[i, j]]

                these_y_coords_metres = storm_object_table[
                    tracking_utils.CENTROID_Y_COLUMN].values[[i, j]]

                if colour_map_object is None:
                    axes_object.plot(these_x_coords_metres,
                                     these_y_coords_metres,
                                     color=this_colour,
                                     linestyle='solid',
                                     linewidth=line_width)
                else:
                    this_point_matrix = numpy.array(
                        [these_x_coords_metres,
                         these_y_coords_metres]).T.reshape(-1, 1, 2)

                    this_segment_matrix = numpy.concatenate(
                        [this_point_matrix[:-1], this_point_matrix[1:]],
                        axis=1)

                    this_time_unix_sec = numpy.mean(storm_object_table[
                        tracking_utils.VALID_TIME_COLUMN].values[[i, j]])

                    this_line_collection_object = LineCollection(
                        this_segment_matrix,
                        cmap=colour_map_object,
                        norm=colour_norm_object)

                    this_line_collection_object.set_array(
                        numpy.array([this_time_unix_sec]))
                    this_line_collection_object.set_linewidth(line_width)
                    axes_object.add_collection(this_line_collection_object)

            these_prev_indices = temporal_tracking.find_immediate_predecessors(
                storm_object_table=storm_object_table, target_row=i)

            # if len(these_prev_indices) > 1:
            #     axes_object.text(
            #         storm_object_table[
            #             tracking_utils.CENTROID_X_COLUMN].values[i],
            #         storm_object_table[
            #             tracking_utils.CENTROID_Y_COLUMN].values[i],
            #         '{0:d}-WAY MERGER'.format(len(these_prev_indices)),
            #         fontsize=12, color='k',
            #         horizontalalignment='left', verticalalignment='top')

            plot_this_start_marker = ((plot_start_markers
                                       and len(these_prev_indices) == 0)
                                      or len(these_object_indices) == 1)

            if plot_this_start_marker:
                if colour_map_object is not None:
                    this_colour = colour_map_object(
                        colour_norm_object(storm_object_table[
                            tracking_utils.VALID_TIME_COLUMN].values[i]))

                if start_marker_type == 'x':
                    this_edge_width = 2
                else:
                    this_edge_width = 1

                axes_object.plot(
                    storm_object_table[
                        tracking_utils.CENTROID_X_COLUMN].values[i],
                    storm_object_table[
                        tracking_utils.CENTROID_Y_COLUMN].values[i],
                    linestyle='None',
                    marker=start_marker_type,
                    markerfacecolor=this_colour,
                    markeredgecolor=this_colour,
                    markersize=start_marker_size,
                    markeredgewidth=this_edge_width)

            plot_this_end_marker = ((plot_end_markers
                                     and len(these_next_indices) == 0)
                                    or len(these_object_indices) == 1)

            if plot_this_end_marker:
                if colour_map_object is not None:
                    this_colour = colour_map_object(
                        colour_norm_object(storm_object_table[
                            tracking_utils.VALID_TIME_COLUMN].values[i]))

                if end_marker_type == 'x':
                    this_edge_width = 2
                else:
                    this_edge_width = 1

                axes_object.plot(
                    storm_object_table[
                        tracking_utils.CENTROID_X_COLUMN].values[i],
                    storm_object_table[
                        tracking_utils.CENTROID_Y_COLUMN].values[i],
                    linestyle='None',
                    marker=end_marker_type,
                    markerfacecolor=this_colour,
                    markeredgecolor=this_colour,
                    markersize=end_marker_size,
                    markeredgewidth=this_edge_width)

    if colour_map_object is None:
        return

    min_plot_latitude_deg = basemap_object.llcrnrlat
    max_plot_latitude_deg = basemap_object.urcrnrlat
    min_plot_longitude_deg = basemap_object.llcrnrlon
    max_plot_longitude_deg = basemap_object.urcrnrlon

    latitude_range_deg = max_plot_latitude_deg - min_plot_latitude_deg
    longitude_range_deg = max_plot_longitude_deg - min_plot_longitude_deg

    if latitude_range_deg > longitude_range_deg:
        orientation_string = 'vertical'
    else:
        orientation_string = 'horizontal'

    colour_bar_object = plotting_utils.plot_linear_colour_bar(
        axes_object_or_matrix=axes_object,
        data_matrix=storm_object_table[
            tracking_utils.VALID_TIME_COLUMN].values,
        colour_map_object=colour_map_object,
        min_value=colour_norm_object.vmin,
        max_value=colour_norm_object.vmax,
        orientation_string=orientation_string,
        extend_min=False,
        extend_max=False,
        fraction_of_axis_length=0.9,
        font_size=COLOUR_BAR_FONT_SIZE)

    if orientation_string == 'horizontal':
        tick_values = colour_bar_object.ax.get_xticks()
    else:
        tick_values = colour_bar_object.ax.get_yticks()

    tick_times_unix_sec = numpy.round(
        colour_norm_object.inverse(tick_values)).astype(int)

    slope_sec_per_sec = (float(last_time_unix_sec - first_time_unix_sec) /
                         (tick_times_unix_sec[-1] - tick_times_unix_sec[0]))

    tick_times_unix_sec = numpy.round(
        first_time_unix_sec + slope_sec_per_sec *
        (tick_times_unix_sec - tick_times_unix_sec[0])).astype(int)

    tick_time_strings = [
        time_conversion.unix_sec_to_string(t, '%Y-%m-%d-%H%M%S')
        for t in tick_times_unix_sec
    ]

    print(tick_time_strings)

    tick_time_strings = [
        time_conversion.unix_sec_to_string(t, COLOUR_BAR_TIME_FORMAT)
        for t in tick_times_unix_sec
    ]

    print(tick_time_strings)

    colour_bar_object.set_ticks(tick_values)
    colour_bar_object.set_ticklabels(tick_time_strings)