sup3r.pipeline.forward_pass.ForwardPass

sup3r.pipeline.forward_pass.ForwardPass#

class ForwardPass(strategy, node_index=0)[source]#

Bases: object

Class to run forward passes on all chunks provided by the given ForwardPassStrategy. The chunks provided by the strategy are all passed through the GAN generator to produce high resolution output.

Initialize ForwardPass with ForwardPassStrategy. The strategy provides the data chunks to run forward passes on

Parameters:

  • strategy (ForwardPassStrategy) – ForwardPassStrategy instance with information on data chunks to run forward passes on.

  • node_index (int) – Index of node used to run forward pass

Methods

get_input_chunk([chunk_index, mode])

Get FowardPassChunk instance for the given chunk index.

get_node_cmd(config)

Get a CLI call to initialize ForwardPassStrategy and run ForwardPass on a single node based on an input config.

pad_source_data(input_data, pad_width, exo_data)

Pad the edges of the source data from the data handler.

run(strategy, node_index)

Runs forward passes on all spatiotemporal chunks for the given node index.

run_chunk(chunk, model_kwargs, model_class, ...)

Run a forward pass on single spatiotemporal chunk.

run_generator(data_chunk, hr_crop_slices, model)

Run forward pass of the generator on smallest data chunk.

Attributes

OUTPUT_HANDLER_CLASS

meta

Meta data dictionary for the forward pass run (to write to output files).
get_input_chunk(chunk_index=0, mode='reflect')[source]#

Get FowardPassChunk instance for the given chunk index.

property meta#

Meta data dictionary for the forward pass run (to write to output files).

pad_source_data(input_data, pad_width, exo_data, mode='reflect')[source]#

Pad the edges of the source data from the data handler.

Parameters:

  • input_data (Union[np.ndarray, da.core.Array]) – Source input data from data handler class, shape is: (spatial_1, spatial_2, temporal, features)

  • pad_width (tuple) – Tuple of tuples with padding width for spatial and temporal dimensions. Each tuple includes the start and end of padding for that dimension. Ordering is spatial_1, spatial_2, temporal.

  • exo_data (dict) – Full exo_handler_kwargs dictionary with all feature entries. See run_generator() for more information.

  • mode (str) – Mode to use for padding. e.g. ‘reflect’.

Returns:

  • out (Union[np.ndarray, da.core.Array]) – Padded copy of source input data from data handler class, shape is: (spatial_1, spatial_2, temporal, features)

  • exo_data (dict) – Same as input dictionary with s_enhance, t_enhance added to each step entry for all features

classmethod run_generator(data_chunk, hr_crop_slices, model, s_enhance=None, t_enhance=None, exo_data=None)[source]#

Run forward pass of the generator on smallest data chunk. Each chunk has a maximum shape given by self.strategy.fwp_chunk_shape.

Parameters:

  • data_chunk (ndarray) – Low res data chunk to go through generator

  • hr_crop_slices (list) – List of slices for extracting cropped region of interest from output. Output can include an extra overlapping boundary to reduce chunking error. The cropping cuts off this padded region before stitching chunks.

  • model (Sup3rGan) – A loaded Sup3rGan model (any model imported from sup3r.models).

  • t_enhance (int) – Factor by which to enhance temporal resolution

  • s_enhance (int) – Factor by which to enhance spatial resolution

  • exo_data (dict | None) – Dictionary of exogenous feature data with entries describing whether features should be combined at input, a mid network layer, or with output. e.g. .. code-block:: JSON

    {
    ‘topography’: {‘steps’: [

    {‘combine_type’: ‘input’, ‘model’: 0, ‘data’: …}, {‘combine_type’: ‘layer’, ‘model’: 0, ‘data’: …}]}

    }

Returns:

ndarray – High resolution data generated by GAN

classmethod get_node_cmd(config)[source]#

Get a CLI call to initialize ForwardPassStrategy and run ForwardPass on a single node based on an input config.

Parameters:

config (dict) – sup3r forward pass config with all necessary args and kwargs to initialize ForwardPassStrategy and run ForwardPass on a single node.

classmethod run(strategy, node_index)[source]#

Runs forward passes on all spatiotemporal chunks for the given node index.

Parameters:

  • strategy (ForwardPassStrategy) – ForwardPassStrategy instance with information on data chunks to run forward passes on.

  • node_index (int) – Index of node on which the forward passes for spatiotemporal chunks will be run.

classmethod run_chunk(chunk: ForwardPassChunk, model_kwargs, model_class, allowed_const, invert_uv=None, meta=None, output_workers=None)[source]#

Run a forward pass on single spatiotemporal chunk.

Parameters:

  • chunk (FowardPassChunk) – Struct with chunk data (including exo data if applicable) and chunk attributes (e.g. chunk specific slices, times, lat/lon, etc)

  • model_kwargs (str | list) – Keyword arguments to send to model_class.load(**model_kwargs) to initialize the GAN. Typically this is just the string path to the model directory, but can be multiple models or arguments for more complex models.

  • model_class (str) – Name of the sup3r model class for the GAN model to load. The default is the basic spatial / spatiotemporal Sup3rGan model. This will be loaded from sup3r.models

  • allowed_const (list | bool) – If your model is allowed to output a constant output, set this to True to allow any constant output or a list of allowed possible constant outputs. See ForwardPassStrategy for more information on this argument.

  • invert_uv (bool) – Whether to convert uv to windspeed and winddirection for writing output. This defaults to True for H5 output and False for NETCDF output.

  • meta (dict | None) – Meta data to write to forward pass output file.

  • output_workers (int | None) – Max number of workers to use for writing forward pass output.

Returns:

  • failed (bool) – Whether the forward pass failed due to constant output.

  • output_data (ndarray) – Array of high-resolution output from generator