mllam / neural-lam

Neural Weather Prediction for Limited Area Modeling
MIT License
64 stars 24 forks source link

Feature Request: Generate GIFs from PNG Images for Each Variable and Level #20

Open sadamov opened 1 month ago

sadamov commented 1 month ago

Description: Currently, the code generates individual PNG images for each variable and level combination at different time steps. To enhance visualization and make it easier to observe changes over time, we propose generating GIF animations from these PNG images.

Proposed Changes:

Iterate over each unique combination of variable name and level. (using #18 )

for var_name, _ in self.selected_vars_units:
    var_indices = self.variable_indices[var_name]
    for lvl_i, _ in enumerate(var_indices):
        # Calculate var_vrange for each index
        lvl = constants.VERTICAL_LEVELS[lvl_i]

        # Get all the images for the current variable and index
        images = sorted(
            glob.glob(
                f"{plot_dir_path}/"
                f"{var_name}_prediction_lvl_{lvl:02}_t_*.png"
            )
        )
        # Generate the GIF
        with imageio.get_writer(
            f"{plot_dir_path}/{var_name}_prediction_lvl_{lvl:02}.gif",
            mode="I",
            fps=1,
        ) as writer:
            for filename in images:
                image = imageio.imread(filename)
                writer.append_data(image)

Benefits:

Considerations:

joeloskarsson commented 1 month ago

I think gif generation is good to have. Saving the pngs and pdfs as well can be useful, but then there are a lot of files. In my plotting scripts I have the formats to save as options, i.e.

    parser.add_argument(
        "--save_png",
        type=int,
        default=1,
        help="If image png files should be saved (default: 1)",
    )
    parser.add_argument(
        "--save_pdf",
        type=int,
        default=0,
        help="If pdf files should be saved (default: 0)",
    )
    parser.add_argument(
        "--save_gif",
        type=int,
        default=1,
        help="If gif should be generated and saved (default: 1)",
    )

A few considerations here (I don't have answers to these, just things to keep in mind when implementing this):

  1. When should such gifs be created? During validation? testing?
  2. Should all plots be sent to wandb, or should some plots only be stored locally? Many of these kinds of files can become quite large.
  3. What plotting should be done in train/test steps and what is done for saved data. It is useful to make some plots during validation to understand model behavior, but for evaluation it is also very much possible to save example forecasts and do plotting separately, by reading those files. Making gifs might be more suitable for this second kind of plotting from saved forecasts. I don't mind if we provide that functionality as well, but then that needs to run from a separate script than train_model.py.