visualize_drawings_experiments.py

"""
visualize_iterative_experiment | Author: Catherine Wong.

Visualizes the results of the iterative experiments for the drawings domain.

Assumes that the directory structure is:
    {experiment_id}/{domain}/{experiment_type}/{replication}/{experiment_type}_{batch_size}


Usage:
python visualize_drawings_experiments.py
    --experiments_id default
    --domains nuts_bolts
    --experiment_type stitch stitch_codex stitch_codex_language_human
    --visualize_codex_results
    --visualize_inventions # Not yet implemented.
"""
import argparse
import json
import os
import random
from collections import defaultdict

import numpy as np
import PIL

import data.drawings.drawings_primitives as drawings_primitives
from data.drawings.grammar import DrawingGrammar
from data.drawings.make_tasks import TASK_DOMAINS  # Drawing task domains.
from dreamcoder.program import Program
from src.models.sample_generator import GPTSampleGenerator

DEFAULT_EXPERIMENTS_DIRECTORY = "experiments_iterative"
DEFAULT_DRAWINGS_GRAMMAR = DrawingGrammar.new_uniform()
DEFAULT_IMAGES_PER_ROW = 10
CODEX_SAMPLE_VISUALIZATION = "codex_sample_visualization.png"
INVENTIONS_VISUALIZTION = "inventions_visualization.png"
DOMAIN_PREFIX = "drawings_"
REPLICATION_PREFIX = "seed_"
CODEX_EXPERIMENT_TAG = "codex"

parser = argparse.ArgumentParser()
parser.add_argument(
    "--experiment_name", required=True, help="Top-level experiment directory."
)
parser.add_argument(
    "--experiment_types",
    nargs="+",
    default=["stitch", "stitch_codex", "stitch_codex_language_human"],
)
parser.add_argument(
    "--domains",
    nargs="+",
    default=TASK_DOMAINS,
)
parser.add_argument(
    "--visualize_codex_results",
    action="store_true",
    help="Visualize programs generated by Codex.",
)
parser.add_argument(
    "--visualize_inventions",
    action="store_true",
    help="Visualize inventions generated by Stitch.",
)


def get_replication_directories(args, domain, experiment_type):
    # Get the directories for this experiment.
    experiment_top_level_directory = os.path.join(
        DEFAULT_EXPERIMENTS_DIRECTORY,
        args.experiment_name,
        DOMAIN_PREFIX + domain,
        experiment_type,
    )
    if not os.path.exists(experiment_top_level_directory):
        print(
            f"Experiment directory not found, skipping visualization: {experiment_top_level_directory}"
        )
        return []

    replication_directories = [
        os.path.join(experiment_top_level_directory, d)
        for d in os.listdir(experiment_top_level_directory)
        if REPLICATION_PREFIX in d
    ]
    return replication_directories


def get_sampled_programs(codex_results, max_to_display=10):
    # Gets sampled programs out of the codex results.
    sampled_programs = [
        p["program"] for p in codex_results["results"]["programs_valid"]
    ]
    sampled_programs = random.sample(
        sampled_programs, min(max_to_display, len(sampled_programs))
    )
    return [Program.parse(string_program) for string_program in sampled_programs]


def get_prompted_programs(codex_results, max_to_display=10):
    # Get programs in the original prompt.
    if "programs" not in codex_results["params"]["body_task_types"]:
        return []

    prompt_programs = []
    for query in codex_results["results_by_query"]:
        for body_task_data in query["prompt"]["body_task_data"]:
            p = body_task_data["task_program"]
            if p is not None:
                prompt_programs.append(body_task_data["task_program"])
        if query["prompt"]["final_task_data"]["task_program"] is not None:
            prompt_programs.append(query["prompt"]["final_task_data"]["task_program"])
    # Randomly sample per prompt.
    prompt_programs = random.sample(
        prompt_programs, min(max_to_display, len(prompt_programs))
    )
    return [Program.parse(string_program) for string_program in prompt_programs]


def stack_images(images):
    min_img_shape = sorted([(np.sum(i.size), i.size) for i in images])[0][1]
    img_merge = np.vstack(
        (np.asarray(i.resize(min_img_shape, PIL.Image.ANTIALIAS)) for i in images)
    )
    img_merge = PIL.Image.fromarray(img_merge)
    return img_merge


def get_blank_spacer(images, percentage=1):
    # Adds a spacer that is percentage * height
    last_image = images[-1]

    spacer = PIL.Image.new(
        "RGB",
        int(last_image.size[0]),
        last_image.size[-1],
        (255, 255, 255),
    )
    return spacer


def visualize_codex_results(args, experiment_type, replication_directory):
    if CODEX_EXPERIMENT_TAG not in experiment_type:
        print(
            f"Assuming {experiment_type} should not have samples visualized, continuing."
        )
        return

    # Get all the batches in this directory, ordered by batch size.
    batch_directories = sorted(
        [d for d in os.listdir(replication_directory) if experiment_type in d],
        key=lambda subdir: int(subdir.split("_")[-1]),
    )
    codex_results_figure = []
    for batch_directory in batch_directories:
        full_batch_directory = os.path.join(replication_directory, batch_directory)
        for iteration in os.listdir(full_batch_directory):
            try:
                # Visualize gpt_query_results.json for each iteration of this batch size.
                codex_results_file = os.path.join(
                    full_batch_directory,
                    iteration,
                    GPTSampleGenerator.query_results_file,
                )
                print(codex_results_file)
                if not os.path.exists(codex_results_file):
                    print(
                        f"Could not find codex_results in {codex_results_file}, continuing."
                    )
                    continue

                with open(codex_results_file) as f:
                    codex_results = json.load(f)

                # Montage the programs in the original prompt.
                prompted_programs = get_prompted_programs(codex_results)
                if len(prompted_programs) > 0:
                    prompted_progams_montage = drawings_primitives.display_programs_as_grid(
                        prompted_programs,
                        color=(0, 0, 0),
                        suptitle=f"Codex prompt program examples: {batch_directory}, stitch_iteration={iteration}",
                        transparent_background=False,
                        ncols=8,
                    )
                    codex_results_figure.append(prompted_progams_montage)

                # Montage the programs sampled from Codex.
                sampled_programs = get_sampled_programs(codex_results)
                if len(sampled_programs) > 0:
                    sampled_programs_montage = drawings_primitives.display_programs_as_grid(
                        sampled_programs,
                        color=(0, 0, 255),
                        suptitle=f"Codex program samples: {batch_directory}, stitch_iteration={iteration}",
                        transparent_background=False,
                        ncols=8,
                    )

                    codex_results_figure.append(sampled_programs_montage)
                # Add a line spacer.
                codex_results_figure.append(get_blank_spacer(codex_results_figure))

            except Exception as e:
                print(e)
                continue

    # Concatenate all the images and save it.
    codex_visualization_output = os.path.join(
        replication_directory, CODEX_SAMPLE_VISUALIZATION
    )
    print(f"Writing out visualization to: {codex_visualization_output}")
    codex_results_image = stack_images(codex_results_figure)
    codex_results_image.save(codex_visualization_output)


def get_iteration_inventions_and_programs(args, iteration, full_batch_directory):
    inventions = set()
    programs = {"codex": set(), "train": set(), "test": set()}

    for split in ["train", "test"]:
        programs_to_rewrite_file = os.path.join(
            full_batch_directory, iteration, split, "stitch_rewrite_input.json"
        )
        programs_rewritten_file = os.path.join(
            full_batch_directory, iteration, split, "stitch_rewrite_output.json"
        )
        with open(programs_to_rewrite_file) as f:
            data = json.load(f)
            split_inventions = [
                Program.parse(p["expression"])
                for p in data["DSL"]["productions"]
                if "#" in p["expression"]
            ]
            inventions.update(split_inventions)
        with open(programs_rewritten_file) as f:
            data = json.load(f)
            for f in data["frontiers"]:
                task_split = split if "codex" not in f["task"] else "codex"
                for program in f["programs"]:
                    programs[task_split].add(Program.parse(program["program"]))
    return inventions, programs


def get_program_primitive_usages(invention, programs):
    usages = defaultdict(list)
    for split in programs:
        for p in programs[split]:
            if str(invention) in p.left_order_tokens():
                usages[split].append(p)
    return usages


def visualize_inventions(args, experiment_type, replication_directory):
    # Get the new inventions at this iteration.
    # Get all the batches in this directory, ordered by batch size.
    batch_directories = sorted(
        [d for d in os.listdir(replication_directory) if experiment_type in d],
        key=lambda subdir: int(subdir.split("_")[-1]),
    )
    for batch_directory in batch_directories:
        full_batch_directory = os.path.join(replication_directory, batch_directory)
        for iteration in os.listdir(full_batch_directory):
            if not iteration.isnumeric():
                continue
            (inventions, programs,) = get_iteration_inventions_and_programs(
                args, iteration, full_batch_directory
            )
            for invention in inventions:
                get_program_primitive_usages(invention, programs)
                import pdb

                pdb.set_trace()


def main(args):
    for domain in args.domains:
        for experiment_type in args.experiment_types:
            replication_directories = get_replication_directories(
                args, domain, experiment_type
            )
            for replication_directory in replication_directories:
                print(
                    f"Now generating visualizations for replication: {replication_directory}"
                )
                if args.visualize_codex_results:
                    visualize_codex_results(
                        args, experiment_type, replication_directory
                    )

                if args.visualize_inventions:
                    visualize_inventions(args, experiment_type, replication_directory)


if __name__ == "__main__":
    args = parser.parse_args()
    main(args)