optimize_farm.py

import json
import sys
from functools import partial
from math import modf
from pathlib import Path
from typing import Annotated, Callable, Final, Literal

import chex
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import typer
from evosax import CMA_ES
from evosax.strategies.cma_es import EvoParams, EvoState
from qdax.core.containers.mapelites_repertoire import (
    MapElitesRepertoire,
    compute_euclidean_centroids,
)
from qdax.core.emitters.cma_improvement_emitter import CMAImprovementEmitter
from qdax.core.emitters.cma_pool_emitter import CMAPoolEmitter
from qdax.core.emitters.mutation_operators import isoline_variation
from qdax.core.emitters.standard_emitters import MixingEmitter
from qdax.core.map_elites import MAPElites
from qdax.types import Centroid
from ribs.visualize import qdax_repertoire_heatmap
from tqdm import trange

from farm import Farm

app = typer.Typer(pretty_exceptions_enable=False)
qd_app = typer.Typer(pretty_exceptions_enable=False)
app.add_typer(qd_app, name="qd", help="Enables `Quality Diversity`.")

PATH_MODE: Final = 0o740


def _init_farm(
    width: chex.Scalar,
    length: chex.Scalar,
    resolution: chex.Scalar,
    cameras: int,
    camera_depth: chex.Scalar,
    horizontal_beam_locations: chex.Array | None,
    vertical_beam_locations: chex.Array | None,
    genotype: Literal["beams", "xy"],
    /,
) -> Farm:
    """Farm initializer helper function.

    Args:
        width: The width of the photo.
        length: The length of the photo.
        resolution: The resolution of the image to convert Cartesian coordinates to
          `image` coordinates. Default is 0.1.
        cameras: The number of total cameras.
        camera_depth: The depth a camera can see from the origin point to the base.
          Default is 3.
        horizontal_beam_locations: The possible horizontal beam locations, that a
          camera can be fixed upon. If None is given the default locations will be used:
          [0m, 3.25m, 6.5m].
        vertical_beam_locations: The possible vertical beam locations, that a camera
          can be fixed upon. If None is given the default locations will be used:
          [0m, 5.125m, 10.25m, 15.375m, 20.5m].
        genotype: The form of the genotype: {'beams', 'xy'}. Default is 'xy'.
            'beams' accepts a genotype with 4 genes: 1) camera orientation, 2) beam
                orientation, 3) beam number, and 4) location on beam. This genotype
                chooses the orientation of the camera, then which beams collection to
                use (vertical or horizontal), which beam to use (from the respective
                collection), and the location of that beam.

            'xy' accepts a genotype with 3 genes: 1) camera orientation, 2) x
                location of camera in image, and 3) y location of camera in image.
                This genotype chooses the orientation of the camera, and its location
                on the image. Then, it finds which of the coordinates is close to on
                of the beams, and clips it to that beam.
        bounds: The bounds for each of the genes. Default is None.
          Example: [[min_1, max_1], [min_2, max_2], [min_3, max_3]].

    Returns:
        A constructed Farm() with the given parameters.
    """

    if horizontal_beam_locations is None:
        horizontal_beam_locations = jnp.linspace(0, length, 3)
    else:
        horizontal_beam_locations = jnp.array(horizontal_beam_locations)

    if vertical_beam_locations is None:
        vertical_beam_locations = jnp.linspace(0, width, 5)
    else:
        vertical_beam_locations = jnp.array(vertical_beam_locations)

    farm = Farm(
        cameras,
        width,
        length,
        horizontal_beam_locations,
        vertical_beam_locations,
        genotype=genotype,
        resolution=resolution,
        triangles_height=camera_depth,
    )

    return farm


def _save_parameters(
    path: Path,
    farm: Farm,
    algorithm: Literal["cmaes", "spread", "spread-xy", "beams_ratio", "beams_ratio"],
    seed: int,
    /,
    *,
    additional_values: dict,
) -> None:
    """Helper function to save the parameters of the experiment as a JSON file.

    Args:
        path: The Path() to save the parameters.
        farm: The Farm() that was created to save.
        algorithm: The algorithm that was chosen.
        seed: The seed that was chosen.
        additional_values: (Optional) Additional values that should be kept.
    """

    with (path / "info.json").open("w") as info:
        args = {
            "algorithm": algorithm,
            "seed": seed,
            "width": farm.width,
            "length": farm.length,
            "resolution": farm.resolution,
            "cameras": farm.cameras,
            "camera_depth": farm.triangles_height,
            "horizontal_beam_locations": farm.horizontal_beam_locations.tolist(),
            "vertical_beam_locations": farm.vertical_beam_locations.tolist(),
            "genotype": farm._genotype_structure,
        }
        args.update(additional_values)
        info.write(json.dumps(args))


def _plot_farm(
    path: Path,
    farm: Farm,
    covered_image: chex.Array,
    camera_xy_points: chex.Array,
    orientations: chex.Array,
    title: str,
    /,
    *,
    filename: str | None = None,
    seperate_coordinates_path: bool = False,
) -> None:
    """Plots the farm with the positioned cameras, and the area that is covered. The
    covered area, is indicated with black pixels. The cameras are drawn as red points,
    with a quiver indicating their orientation. The reasulting files have a PDF, and
    a CSV format respectively.

    Args:
        path: The Path() to save the farm plot.
        farm: The Farm() to be plotted.
        covered_image: The mapping of the covered pixels.
        camera_xy_points: The points of the cameras.
        orientations: The orientation os the cameras.
        title: The title of the plot.
        filename: (Optional) The desired filename for the file.
        seperate_coordinates_path: (Optional) If 'True', the coordinates of the
          cameras + orientation will be saved in a different path, else if 'False',
          the file will be saved in the same path as the PDF.
    """

    fig, ax = plt.subplots(1, 1, figsize=(11.693, 8.268), layout="constrained")
    ax.imshow(
        1 - covered_image,
        "gray",
        origin="lower",
        extent=[0, farm.width, 0, farm.length],
    )
    ax.scatter(camera_xy_points.at[0].get(), camera_xy_points.at[1].get(), 10, "red")

    for h_beam in farm.horizontal_beam_locations:
        plt.axhline(h_beam, color="g")

    for v_beam in farm.vertical_beam_locations:
        plt.axvline(v_beam, color="g")

    for x, y, orientation in zip(
        camera_xy_points.at[0].get(), camera_xy_points.at[1].get(), orientations
    ):
        orientation = jnp.radians(orientation)
        ax.quiver(
            x,
            y,
            farm.triangles_height * jnp.cos(orientation),
            farm.triangles_height * jnp.sin(orientation),
            units="width",
            color="r",
        )

    ax.set_title(title)
    ax.set_xlim(-0.5, farm.width + 0.5)
    ax.set_ylim(-0.5, farm.length + 0.5)
    ax.set_xlabel("Width")
    ax.set_ylabel("Length")
    if filename is None:
        filename = "farm_solution.pdf"

    fig.savefig(path / filename)

    plt.close()

    coordinates_filename: str = "coordinates.csv"
    if filename is not None:
        coordinates_filename = f"{filename.removesuffix('.pdf')}.csv"

    if seperate_coordinates_path:
        coordinates_path: Path = path / "coordinates"
        coordinates_path.mkdir(PATH_MODE, exist_ok=True)
    else:
        coordinates_path = path

    with (coordinates_path / coordinates_filename).open("w") as file:
        file.write("camera_no,x,y,orientation\n")
        for index in range(orientations.size):
            file.write(
                f"{index + 1},{camera_xy_points[0, index]},"
                f"{camera_xy_points[1, index]},{orientations[index]}\n"
            )


def _plot_all_farms(
    path: Path,
    farm: Farm,
    repertoire: MapElitesRepertoire,
    total_generations: int,
    batch_size: int,
    emitters: int,
    /,
    *,
    descriptor_names: list[str] = None,
) -> None:
    """Plots multiple farms from a MapElitesRepertoire utilizing _plot_farm().

    Args:
        path: The Path() to save the farm plot.
        farm: The Farm() to be plotted.
        repertoire: The MapElitesRepertoire() that contains all the behaviour
          descriptors that should be plotted.
        total_generations: The total training evaluations generations to be included in
          the title.
        batch_size: The batch size that was used for each emitter to be included in the
          title.
        emitters: The amount of emitters that were used to be included in the title.
        descriptor_names: (Optional) A list containing all the names of the descriptors.
          Serves as a helper when plotting "used-beams", and "beams_ratio".
    """

    non_empty_fitness_indices = repertoire.fitnesses != -jnp.inf

    solutions_path: Path = path / "solutions"
    solutions_path.mkdir(PATH_MODE, exist_ok=True)

    fitnesses, covered_images, _, origin_points, orientations = jax.vmap(
        farm.fitness, (0)
    )(repertoire.genotypes[non_empty_fitness_indices])

    solutions_listed = {}

    for current_solution_index, true_index in enumerate(
        jnp.argwhere(non_empty_fitness_indices).flatten().tolist()
    ):
        filename: str = ""
        if descriptor_names:
            for index, descriptor_name in enumerate(descriptor_names):
                filename += (
                    f"{descriptor_name}_{repertoire.descriptors[true_index, index]}"
                )

            filename = filename.removesuffix("-") + ".pdf"

        else:
            filename = f"{repertoire.centroids[true_index, :]}.pdf"

        solutions_listed[f"{filename}"] = repertoire.fitnesses[true_index].tolist()

        _plot_farm(
            solutions_path,
            farm,
            covered_images.at[current_solution_index].get(),
            origin_points.at[current_solution_index].get(),
            orientations.at[current_solution_index].get(),
            f"Coverage: {fitnesses.at[current_solution_index].get()}, "
            f"{total_generations} generations, batch size: {batch_size}, "
            f"emitters: {emitters}",
            filename=filename,
            seperate_coordinates_path=True,
        )

    solutions_listed = dict(
        sorted(solutions_listed.items(), key=lambda item: item[1], reverse=True)
    )

    with (path / "solutions_info.json").open("w") as file:
        file.write(json.dumps(solutions_listed))


def _plot_behaviour_space_genotype_space(
    path: Path,
    repertoire: MapElitesRepertoire,
    total_generations: int,
    batch_size: int,
    emitters: int,
    vmin: chex.Scalar,
    vmax: chex.Scalar,
    /,
    *,
    x_label: str | None = None,
    y_label: str | None = None,
) -> None:
    """Behaviour space plot helper for spead and spead-xy.

    Args:
        path: The Path() to save the farm plot.
        repertoire: The MapElitesRepertoire() that contains all the solutions that
          should be plotted.
        total_generations: The total training evaluations generations to be included in
        the title.
        batch_size: The batch size that was used for each emitter to be included in the
        title.
        emitters: The amount of emitters that were used to be included in the title.
        vmin: The minimum fitness that was found when filing the repertoire.
        vmax: The maximum fitness that was found when filing the repertoire.
        x_label: (Optional) The label that should be used for the X-axis.
        y_label: (Optional) The label that should be used for the Y-axis.
    """

    fig, ax = plt.subplots(1, 1, figsize=(11.693, 8.268), layout="constrained")
    qdax_repertoire_heatmap(
        repertoire,
        jnp.full(
            (repertoire.descriptors.shape[1], 2),
            jnp.array([0, 1]),
        ),
        ax=ax,
        cmap="plasma",
        ec="#00000030",
        vmin=vmin,
        vmax=vmax,
        cbar_kwargs={"location": "bottom", "shrink": 0.4},
    )

    ax.set_title(
        f"Behaviour space after {total_generations} generations, using batch size "
        f"{batch_size} and {emitters} emitters. Total evaluations: "
        f"{total_generations * batch_size * emitters:,}"
    )
    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)

    fig.savefig(path / "farm_qd_space.pdf")

    plt.close()


def _save_results(
    path: Path, results: dict, /, *, metrics: dict[str, chex.Array] | None = None
) -> None:
    """Helps save the results as a JSON file.

    Args:
        path: The Path() to save the results.
        results: The results to be saved.
        metrics: (Optional) Metrics of the results that should also be saved. Mainly to
          be used for QD metrics.
    """

    with (path / "results.json").open("w") as file:
        file.write(json.dumps(results))

    if metrics:
        with (path / "metrics.json").open("w") as file:
            for key, value in metrics.items():
                metrics[key] = value.tolist()[0]

            file.write(json.dumps(metrics))


def cmaes_optimization(
    generations: int,
    seed: int,
    /,
    *,
    population: int | None = None,
    width: chex.Scalar = 20.5,
    length: chex.Scalar = 6.5,
    resolution: chex.Scalar = 0.1,
    cameras: int = 6,
    camera_depth: chex.Scalar = 5,
    horizontal_beam_locations: chex.Array | None = None,
    vertical_beam_locations: chex.Array | None = None,
    genotype: Literal["beams", "xy"] = "xy",
    custom_path_prefix: Path | None = None,
) -> None:
    """Performs optimization on the given Farm specifications using CMA-ES.

    Args:
        generations: The maximum total number of evaluation generations.
        seed: The seed to be used in JAX's PRNG.
        population: (Optional) The population of each generation (the number of
          evaulations per generation). If None, the default population will be
          calculated using:
          4 + floor(3 * ln(farm.dims)). Default is None.
        width: (Optional) The width of the photo. Default is 20.5.
        length: (Optional) The length of the photo. Default is 6.5.
        resolution: (Optional) The resolution of the image to convert Cartesian
          coordinates to `image` coordinates. Default is 0.1.
        cameras: (Optional) The number of total cameras. Default is 6.
        camera_depth: (Optional) The depth a camera can see from the origin point to
          the base. Default is 5.
        horizontal_beam_locations: (Optional) The possible horizontal beam locations,
          that a camera can be fixed upon. If None is given the default locations will
          be used: [0m, 3.25m, 6.5m]. Default is None.
        vertical_beam_locations: (Optional) The possible vertical beam locations, that
          a camera can be fixed upon. If None is given the default locations will be
          used: [0m, 5.125m, 10.25m, 15.375m, 20.5m]. Default is None.
        genotype: (Optional) The form of the genotype: {'beams', 'xy'}. Default is 'xy'.
            'beams' accepts a genotype with 4 genes: 1) camera orientation, 2) beam
                orientation, 3) beam number, and 4) location on beam. This genotype
                chooses the orientation of the camera, then which beams collection to
                use (vertical or horizontal), which beam to use (from the respective
                collection), and the location of that beam.

            'xy' accepts a genotype with 3 genes: 1) camera orientation, 2) x
                location of camera in image, and 3) y location of camera in image.
                This genotype chooses the orientation of the camera, and its location
                on the image. Then, it finds which of the coordinates is close to on
                of the beams, and clips it to that beam.
        custom_path_prefix: (Optional) A prefix of the desired location to save all the
          produced files by the function. If default is given, the current location will
          be used (execution location). Default is None.
    """

    farm: Farm = _init_farm(
        width,
        length,
        resolution,
        cameras,
        camera_depth,
        horizontal_beam_locations,
        vertical_beam_locations,
        genotype,
    )

    if population is None:
        population = (4 + jnp.floor(3 * jnp.log(farm.dims))).astype(int)

    key: chex.PRNGKey = jax.random.key(seed)
    strategy = CMA_ES(population, farm.dims)
    es_params: EvoParams = strategy.default_params.replace(clip_min=0, clip_max=1)
    strategy_state: EvoState = strategy.initialize(key, es_params)

    path: Path
    if custom_path_prefix:
        path = Path(f"{custom_path_prefix}/experiment-seed_{seed}")
    else:
        path = Path(f"experiment-seed_{seed}")

    path.mkdir(PATH_MODE, exist_ok=True, parents=True)
    _save_parameters(
        path,
        farm,
        "cmaes",
        seed,
        additional_values={"generations": generations, "population": int(population)},
    )

    results: dict = {}

    for generation in trange(1, generations + 1, desc="Generations", file=sys.stdout):
        key, rng_gen, rng_eval = jax.random.split(key, 3)
        x, strategy_state = strategy.ask(rng_gen, strategy_state, es_params)
        fitness, _, _, _, _ = jax.vmap(farm.fitness, (0))(x)
        fitness = 1 - fitness
        strategy_state = strategy.tell(x, fitness, strategy_state, es_params)
        results[generation] = {
            "solution": strategy_state.best_member.tolist(),
            "fitness": 1 - float(strategy_state.best_fitness),
        }
        sys.stdout.flush()

    _save_results(path, results)

    fitness, covered_image, _, points, orientations = farm.fitness(
        strategy_state.best_member
    )

    _plot_farm(
        path,
        farm,
        covered_image,
        points,
        orientations,
        f"Coverage: {fitness}, {generations} generations, population: {population}",
    )


def _qd_helper(
    generations: int,
    operator_type: Literal["cma-me", "isolinedd"],
    batch_size: int,
    emitters: int,
    farm: Farm,
    centroids: Centroid,
    simulation_function: Callable[[chex.Array], tuple[chex.Scalar, chex.Array, dict]],
    key: chex.PRNGKey,
    /,
    *,
    total_logs: int | None = None,
) -> tuple[
    dict[dict[str, list]],
    MapElitesRepertoire,
    dict[str, chex.Array],
    chex.Scalar,
    chex.Scalar,
    chex.PRNGKey,
]:
    """QD helper function.

    Args:
        generations: The maximum total number of evaluation generations.
        operator_type: The operator algorithm to be used for finding solutions.
        batch_size: The total solutions to be found at each generaton.
        emitters: The total emitters to be used for finding solutions. Each emitters
          can find 'batch_size' solutions each generation. Only effective when
          operator_type='cma-me'.
        farm: The Farm() that we want to find optimal solutions.
        centroids: The Centroids of the behaviour space.
        simulation_function: The function that will provide the fitness, and the
          behaviour descriptors of the given solution (genotype).
        key: The chex.PRNGKey that should be used for the experiment.
        total_logs: (Optional) The number of total logs per generation that should be
          taken. If None is given, only the last result will be taken. Default is
          None. Note that as the number of logs increases, so does the completion
          time.

    Returns:
        A dictionary with the results taken after the total_logs, the repertoire with
        the found solutions and the behaviour descriptors, the metrics of the
        repertoire (given total_logs), the minimum fitness, the maximum fitness, and
        the updated PRNGKey that was provided.
    """

    match operator_type:
        case "cma-me":
            emitter = CMAImprovementEmitter(batch_size, farm.dims, centroids, 0.5)
            emitter = CMAPoolEmitter(emitters, emitter)
        case "isolinedd":
            emitter = MixingEmitter(
                None,
                partial(isoline_variation, iso_sigma=0.1, line_sigma=0.2),
                1,
                batch_size,
            )
        case _:
            raise ValueError(
                f"Provided opreator does not exist. operator = '{operator_type}'."
            )

    initial_population: chex.Array = jnp.empty((batch_size, farm.dims))

    @jax.jit
    def _scoring_function(Xs: chex.Array, key: chex.PRNGKey):
        fitnesses, descriptors, extras = jax.vmap(jax.jit(simulation_function), (0))(Xs)
        return fitnesses, descriptors, extras, key

    @jax.jit
    def _metrics_function(repertoire: MapElitesRepertoire) -> dict[str, chex.Array]:
        not_empty_grid: chex.Array = repertoire.fitnesses != -jnp.inf
        qd_score: chex.Array = repertoire.fitnesses.sum(where=not_empty_grid)
        mean_fitness: chex.Array = repertoire.fitnesses.mean(where=not_empty_grid)
        total_elites: chex.Array = jnp.count_nonzero(not_empty_grid)
        coverage: chex.Array = not_empty_grid.mean()
        max_fitness: chex.Array = repertoire.fitnesses.max()
        max_fitness_index: chex.Array = repertoire.fitnesses.argmax()

        return {
            "total_elites": total_elites,
            "coverage": coverage,
            "qd_score": qd_score,
            "mean_fitness": mean_fitness,
            "max_fitness": max_fitness,
            "max_fitness_descriptors": repertoire.descriptors[max_fitness_index],
            "max_fitness_centroids": repertoire.centroids[max_fitness_index],
        }

    map_elites = MAPElites(_scoring_function, emitter, _metrics_function)

    repertoire: MapElitesRepertoire

    repertoire, emitter_state, key = map_elites.init(initial_population, centroids, key)

    results: dict = {}

    vmin: float = jnp.nan
    vmax: float = jnp.nan

    if total_logs is None:
        total_logs = 1

    loop_size_fractional, loop_size_integer = modf(generations / total_logs)

    loop_size_fractional = round(loop_size_fractional, 2)
    loop_size_integer = int(loop_size_integer)

    total_fractional = 0.0

    for generation in trange(total_logs, desc="Generations", file=sys.stdout):
        total_fractional += loop_size_fractional
        total_fractional, loop_extension = modf(total_fractional)
        loop_size_total = loop_size_integer + int(loop_extension)

        (repertoire, emitter_state, key), metrics = jax.lax.scan(
            map_elites.scan_update,
            (repertoire, emitter_state, key),
            (),
            loop_size_total,
        )

        results[generation] = {
            "solutions": repertoire.genotypes.tolist(),
            "descriptors": repertoire.descriptors.tolist(),
            "fitnesses": repertoire.fitnesses.tolist(),
        }

        vmin = jnp.nanmin(
            jnp.array(
                [vmin, repertoire.fitnesses[repertoire.fitnesses != -jnp.inf].min()]
            )
        )
        vmax = jnp.nanmax(
            jnp.array(
                [vmax, repertoire.fitnesses[repertoire.fitnesses != jnp.inf].max()]
            )
        )

        sys.stdout.flush()

    return results, repertoire, metrics, vmin, vmax, key


def qd_spread_optimization(
    generations: int,
    seed: int,
    operator_type: Literal["cma-me", "isolinedd"],
    /,
    *,
    batch_size: int = 36,
    emitters: int = 15,
    grid_size: int = 1000,
    only_xy: bool = False,
    width: chex.Scalar = 20.5,
    length: chex.Scalar = 6.5,
    resolution: chex.Scalar = 0.1,
    cameras: int = 6,
    camera_depth: chex.Scalar = 5,
    horizontal_beam_locations: chex.Array | None = None,
    vertical_beam_locations: chex.Array | None = None,
    genotype: Literal["beams", "xy"] = "xy",
    total_logs: int | None = None,
    custom_path_prefix: Path | None = None,
) -> None:
    """Performs optimization on the given Farm specifications using QD + the 'spread'
    or the 'spread-xy' behaviour function. 'spread' uses X, Y, and orientation to
    calculate the distance, and 'spread-xy' uses only X, and Y.

    Args:
        generations: The maximum total number of evaluation generations.
        seed: The seed to be used in JAX's PRNG.
        operator_type: The operator algorithm to be used for finding solutions.
        batch_size: (Optional) The total solutions to be found at each generaton.
        emitters: (Optional) The total emitters to be used for finding solutions. Each
          emitters can find 'batch_size' solutions each generation. Only effective when
          operator_type='cma-me'.
        grid_size: (Optional) The grid size to be used for the behaviour descriptor.
        only_xy: (Optional) If True, the 'spread-xy' descriptor will be used, else if
          False, the 'spead' descriptor will be used. Default is False.
        width: (Optional) The width of the photo. Default is 20.5.
        length: (Optional) The length of the photo. Default is 6.5.
        resolution: (Optional) The resolution of the image to convert Cartesian
          coordinates to `image` coordinates. Default is 0.1.
        cameras: (Optional) The number of total cameras. Default is 6.
        camera_depth: (Optional) The depth a camera can see from the origin point to
          the base. Default is 5.
        horizontal_beam_locations: (Optional) The possible horizontal beam locations,
          that a camera can be fixed upon. If None is given the default locations will
          be used: [0m, 3.25m, 6.5m]. Default is None.
        vertical_beam_locations: (Optional) The possible vertical beam locations, that
          a camera can be fixed upon. If None is given the default locations will be
          used: [0m, 5.125m, 10.25m, 15.375m, 20.5m]. Default is None.
        genotype: (Optional) The form of the genotype: {'beams', 'xy'}. Default is 'xy'.
            'beams' accepts a genotype with 4 genes: 1) camera orientation, 2) beam
                orientation, 3) beam number, and 4) location on beam. This genotype
                chooses the orientation of the camera, then which beams collection to
                use (vertical or horizontal), which beam to use (from the respective
                collection), and the location of that beam.

            'xy' accepts a genotype with 3 genes: 1) camera orientation, 2) x
                location of camera in image, and 3) y location of camera in image.
                This genotype chooses the orientation of the camera, and its location
                on the image. Then, it finds which of the coordinates is close to on
                of the beams, and clips it to that beam.
        custom_path_prefix: (Optional) A prefix of the desired location to save all the
          produced files by the function. If default is given, the current location will
          be used (execution location). Default is None.
    """

    farm: Farm = _init_farm(
        width,
        length,
        resolution,
        cameras,
        camera_depth,
        horizontal_beam_locations,
        vertical_beam_locations,
        genotype,
    )

    if operator_type != "cma-me":
        emitters = 1

    def _simulate(X: chex.Array) -> tuple[chex.Scalar, chex.Array, dict]:
        fitness, _, _, _, _ = farm.fitness(X)
        return (
            fitness,
            jnp.array([farm.behaviour_descriptor_spread(X, only_xy=only_xy)]),
            {},
        )

    key: chex.PRNGKey = jax.random.key(seed)

    path: Path
    if custom_path_prefix:
        path = Path(f"{custom_path_prefix}/experiment-seed_{seed}")
    else:
        path = Path(f"experiment-seed_{seed}")

    path.mkdir(PATH_MODE, exist_ok=True, parents=True)
    _save_parameters(
        path,
        farm,
        "spread" if not only_xy else "spread-xy",
        seed,
        additional_values={
            "generations": generations,
            "batch_size": batch_size,
            "operator_type": operator_type,
            "emitters": emitters,
            "grid_size": grid_size,
        },
    )

    centroids = compute_euclidean_centroids([grid_size], 0, 1)

    results, repertoire, metrics, vmin, vmax, key = _qd_helper(
        generations,
        operator_type,
        batch_size,
        emitters,
        farm,
        centroids,
        _simulate,
        key,
        total_logs=total_logs,
    )

    _save_results(path, results, metrics=metrics)

    _plot_behaviour_space_genotype_space(
        path,
        repertoire,
        generations,
        batch_size,
        emitters,
        vmin,
        vmax,
        x_label="Spread only X,Y",
    )

    _plot_all_farms(
        path,
        farm,
        repertoire,
        generations,
        batch_size,
        emitters,
    )


def qd_beams_optimization(
    generations: int,
    seed: int,
    operator_type: Literal["cma-me", "isolinedd"],
    behaviour_descriptor: Literal["used-beams", "beams-ratio"],
    /,
    *,
    batch_size: int = 36,
    emitters: int = 15,
    width: chex.Scalar = 20.5,
    length: chex.Scalar = 6.5,
    resolution: chex.Scalar = 0.1,
    cameras: int = 6,
    camera_depth: chex.Scalar = 5,
    horizontal_beam_locations: chex.Array | None = None,
    vertical_beam_locations: chex.Array | None = None,
    genotype: Literal["beams", "xy"] = "xy",
    total_logs: int | None = None,
    custom_path_prefix: Path | None = None,
) -> None:
    """Performs optimization on the given Farm specifications using QD + the
    'used-beams' or the 'beams_ratio' behaviour function. 'used-beams' finds solutions
    where some beams are used, whilst some other are not; 'beams_ratio', finds solutions
    where each beam is being used from-and-up-to [0 - 0.334], [0.334 - 0.667],
    and [0.667, 1].

    Args:
        generations: The maximum total number of evaluation generations.
        seed: The seed to be used in JAX's PRNG.
        operator_type: The operator algorithm to be used for finding solutions.
        behaviour_descriptor: The descriptor that should be used.
        batch_size: (Optional) The total solutions to be found at each generaton.
        emitters: (Optional) The total emitters to be used for finding solutions. Each
          emitters can find 'batch_size' solutions each generation. Only effective when
          operator_type='cma-me'.
        grid_size: (Optional) The grid size to be used for the behaviour descriptor.
        only_xy: (Optional) If True, the 'spread-xy' descriptor will be used, else if
          False, the 'spead' descriptor will be used. Default is False.
        width: (Optional) The width of the photo. Default is 20.5.
        length: (Optional) The length of the photo. Default is 6.5.
        resolution: (Optional) The resolution of the image to convert Cartesian
          coordinates to `image` coordinates. Default is 0.1.
        cameras: (Optional) The number of total cameras. Default is 6.
        camera_depth: (Optional) The depth a camera can see from the origin point to
          the base. Default is 5.
        horizontal_beam_locations: (Optional) The possible horizontal beam locations,
          that a camera can be fixed upon. If None is given the default locations will
          be used: [0m, 3.25m, 6.5m]. Default is None.
        vertical_beam_locations: (Optional) The possible vertical beam locations, that
          a camera can be fixed upon. If None is given the default locations will be
          used: [0m, 5.125m, 10.25m, 15.375m, 20.5m]. Default is None.
        genotype: (Optional) The form of the genotype: {'beams', 'xy'}. Default is 'xy'.
            'beams' accepts a genotype with 4 genes: 1) camera orientation, 2) beam
                orientation, 3) beam number, and 4) location on beam. This genotype
                chooses the orientation of the camera, then which beams collection to
                use (vertical or horizontal), which beam to use (from the respective
                collection), and the location of that beam.

            'xy' accepts a genotype with 3 genes: 1) camera orientation, 2) x
                location of camera in image, and 3) y location of camera in image.
                This genotype chooses the orientation of the camera, and its location
                on the image. Then, it finds which of the coordinates is close to on
                of the beams, and clips it to that beam.
        custom_path_prefix: (Optional) A prefix of the desired location to save all the
          produced files by the function. If default is given, the current location will
          be used (execution location). Default is None.
    """

    farm: Farm = _init_farm(
        width,
        length,
        resolution,
        cameras,
        camera_depth,
        horizontal_beam_locations,
        vertical_beam_locations,
        genotype,
    )

    if operator_type != "cma-me":
        emitters = 1

    def _simulate_used(X: chex.Array) -> tuple[chex.Scalar, chex.Array, dict]:
        fitness, _, _, camera_points, _ = farm.fitness(X)
        beams: dict = farm.behaviour_descriptor_used_beams(camera_points=camera_points)

        return fitness, jnp.array(list(beams.values())), {}

    def _simulate_ratio(X: chex.Array) -> tuple[chex.Scalar, chex.Array, dict]:
        fitness, _, _, camera_points, _ = farm.fitness(X)
        beams: dict = farm.behaviour_descriptor_beams_ratio(camera_points=camera_points)

        return fitness, jnp.array(list(beams.values())), {}

    key: chex.PRNGKey = jax.random.key(seed)

    grid_size: int
    simulate_fn: Callable[[chex.Array], tuple[chex.Scalar, chex.Array, dict]]

    match behaviour_descriptor:
        case "used-beams":
            grid_size = 2
            simulate_fn = _simulate_used
        case "beams-ratio":
            grid_size = 3
            simulate_fn = _simulate_ratio
        case _:
            raise ValueError(
                f"Behaviour descriptor '{behaviour_descriptor}' does not exist."
            )

    path: Path
    if custom_path_prefix:
        path = Path(f"{custom_path_prefix}/experiment-seed_{seed}")
    else:
        path = Path(f"experiment-seed_{seed}")

    path.mkdir(PATH_MODE, exist_ok=True, parents=True)
    _save_parameters(
        path,
        farm,
        behaviour_descriptor,
        seed,
        additional_values={
            "generations": generations,
            "batch_size": batch_size,
            "operator_type": operator_type,
            "emitters": emitters,
        },
    )

    centroids: Centroid = compute_euclidean_centroids(
        [grid_size]
        * (farm.horizontal_beam_locations.size + farm.vertical_beam_locations.size),
        0,
        1,
    )

    results, repertoire, metrics, vmin, vmax, key = _qd_helper(
        generations,
        operator_type,
        batch_size,
        emitters,
        farm,
        centroids,
        simulate_fn,
        key,
        total_logs=total_logs,
    )

    descriptor_names = list(
        farm.behaviour_descriptor_used_beams(genotypes=repertoire.genotypes[0]).keys()
    )

    _save_results(path, results, metrics=metrics)

    _plot_all_farms(
        path,
        farm,
        repertoire,
        generations,
        batch_size,
        emitters,
        descriptor_names=descriptor_names,
    )


@qd_app.command(
    "spread",
    help="Uses spread, as the descriptor function ranging between [0, 1]; 0 means "
    "overlapping (not spread out), 1 means scattered (not overlapping, as distant as "
    "they can be).",
)
def spread(
    generations: Annotated[
        int,
        typer.Argument(help="The number of total generations to perform CMA-ES."),
    ],
    seed: Annotated[
        int,
        typer.Argument(help="The seed for JAX's PRNG."),
    ],
    operator_type: Annotated[
        str,
        typer.Argument(
            help="Indicates the QD variation operator. The operators available are: "
            "'cma-me', 'isolinedd'."
        ),
    ],
    batch_size: Annotated[
        int,
        typer.Option(
            "--batch",
            "-b",
            help="The total batch size of CMA-ME. Default is 36 batch size.",
        ),
    ] = 36,
    emitters: Annotated[
        int,
        typer.Option(
            "--emitters",
            "-e",
            help="The total emitters to use with CMA-ME. Default is 15 emitters.",
        ),
    ] = 15,
    grid_size: Annotated[
        int,
        typer.Option(
            "--grid",
            "-g",
            help="The total cells to keep diversity in the behaviour space."
            " Default is 1000 cells",
        ),
    ] = 1000,
    only_xy: Annotated[
        bool,
        typer.Option(
            "--only_xy",
            "-xy",
            help="Forces the spread operator to use only X and Y for the spread "
            "calculation. Default is using X, Y and orientation",
        ),
    ] = False,
    width: Annotated[
        float,
        typer.Option(
            "--width",
            "-w",
            help="The width of the farm in meters. Default is 20.5 meters.",
        ),
    ] = 20.5,
    length: Annotated[
        float,
        typer.Option(
            "--length",
            "-l",
            help="The length of the farm in meters. Default is 6.5 meters.",
        ),
    ] = 6.5,
    resolution: Annotated[
        float,
        typer.Option(
            "--resolution",
            "-r",
            help="The resolution of the image in meters. Default is 0.1 meters (10cm).",
        ),
    ] = 0.1,
    cameras: Annotated[
        int,
        typer.Option(
            "--cameras",
            "-c",
            help="The number of cameras to be placed in the farm. Default is 6.",
        ),
    ] = 6,
    camera_depth: Annotated[
        float,
        typer.Option(
            "--camera_depth",
            "-d",
            help="The depth of the camera in meters; i.e. how far can a camera should "
            "see clearly. Default is 5 meters.",
        ),
    ] = 5,
    horizontal_beam_locations: Annotated[
        list[float],
        typer.Option(
            "--horizontal_beam",
            "-hb",
            help="The horizontal beam locations that a camera can be placed on. For "
            "many locations use '-hb' multiple times. Default behaviour will create 3 "
            "beams with equal interval between them.",
        ),
    ] = None,
    vertical_beam_locations: Annotated[
        list[float],
        typer.Option(
            "--vertical_beam",
            "-vb",
            help="The vertical beam locations that a camera can be placed on. For "
            "many locations use '-vb' multiple times. Default behaviour will create 5 "
            "beams with equal interval between them.",
        ),
    ] = None,
    genotype: Annotated[
        str,
        typer.Option(
            "--genotype",
            "-gt",
            help="The genotype to be used Available genotypes: ['beams', 'xy']. "
            "Default is 'xy'.",
        ),
    ] = "xy",
    total_logs: Annotated[
        int,
        typer.Option(
            "--total_logs",
            help="The number of total logs that should be taken with a similar "
            "interval. Default is None; every generation is logged.",
        ),
    ] = None,
) -> None:
    qd_spread_optimization(
        generations,
        seed,
        operator_type,
        batch_size=batch_size,
        emitters=emitters,
        grid_size=grid_size,
        only_xy=only_xy,
        width=width,
        length=length,
        resolution=resolution,
        cameras=cameras,
        camera_depth=camera_depth,
        horizontal_beam_locations=horizontal_beam_locations,
        vertical_beam_locations=vertical_beam_locations,
        genotype=genotype,
        total_logs=total_logs,
    )


@qd_app.command(
    "beams",
    help="Invokes the beam-oriented descriptor functions, where the is a descriptor "
    "for each given beams; both horizontal and vertical.",
)
def beams(
    generations: Annotated[
        int,
        typer.Argument(help="The number of total generations to perform CMA-ES."),
    ],
    seed: Annotated[
        int,
        typer.Argument(help="The seed for JAX's PRNG."),
    ],
    operator_type: Annotated[
        str,
        typer.Argument(
            help="Indicates the QD variation operator. The operators available are: "
            "'cma-me', 'isolinedd'."
        ),
    ],
    behaviour_descriptor: Annotated[
        str,
        typer.Argument(
            help="Beams behaviour function to use. Available descriptors: "
            "'used-beams', 'beams-ratio'. The 'used-beams' decsriptor finds solutions "
            "where some beams are used, whilst some other are not; the 'beams-ratio', "
            "which finds solutions where each beam is being used from-and-up-to "
            "[0 - 0.334], [0.334 - 0.667], and [0.667, 1].",
        ),
    ],
    batch_size: Annotated[
        int,
        typer.Option(
            "--batch",
            "-b",
            help="The total batch size of CMA-ME. Default is 36 batch size.",
        ),
    ] = 36,
    emitters: Annotated[
        int,
        typer.Option(
            "--emitters",
            "-e",
            help="The total emitters to use with CMA-ME. Default is 15 emitters.",
        ),
    ] = 15,
    width: Annotated[
        float,
        typer.Option(
            "--width",
            "-w",
            help="The width of the farm in meters. Default is 20.5 meters.",
        ),
    ] = 20.5,
    length: Annotated[
        float,
        typer.Option(
            "--length",
            "-l",
            help="The length of the farm in meters. Default is 6.5 meters.",
        ),
    ] = 6.5,
    resolution: Annotated[
        float,
        typer.Option(
            "--resolution",
            "-r",
            help="The resolution of the image in meters. Default is 0.1 meters (10cm).",
        ),
    ] = 0.1,
    cameras: Annotated[
        int,
        typer.Option(
            "--cameras",
            "-c",
            help="The number of cameras to be placed in the farm. Default is 6.",
        ),
    ] = 6,
    camera_depth: Annotated[
        float,
        typer.Option(
            "--camera_depth",
            "-d",
            help="The depth of the camera in meters; i.e. how far can a camera should "
            "see clearly. Default is 5 meters.",
        ),
    ] = 5,
    horizontal_beam_locations: Annotated[
        list[float],
        typer.Option(
            "--horizontal_beam",
            "-hb",
            help="The horizontal beam locations that a camera can be placed on. For "
            "many locations use '-hb' multiple times. Default behaviour will create 3 "
            "beams with equal interval between them.",
        ),
    ] = None,
    vertical_beam_locations: Annotated[
        list[float],
        typer.Option(
            "--vertical_beam",
            "-vb",
            help="The vertical beam locations that a camera can be placed on. For "
            "many locations use '-vb' multiple times. Default behaviour will create 5 "
            "beams with equal interval between them.",
        ),
    ] = None,
    genotype: Annotated[
        str,
        typer.Option(
            "--genotype",
            "-gt",
            help="The genotype to be used Available genotypes: ['beams', 'xy']. "
            "Default is 'xy'.",
        ),
    ] = "xy",
    total_logs: Annotated[
        int,
        typer.Option(
            "--total_logs",
            help="The number of total logs that should be taken with a similar "
            "interval. Default is None; every generation is logged.",
        ),
    ] = None,
) -> None:
    qd_beams_optimization(
        generations,
        seed,
        operator_type,
        behaviour_descriptor,
        batch_size=batch_size,
        emitters=emitters,
        width=width,
        length=length,
        resolution=resolution,
        cameras=cameras,
        camera_depth=camera_depth,
        horizontal_beam_locations=horizontal_beam_locations,
        vertical_beam_locations=vertical_beam_locations,
        genotype=genotype,
        total_logs=total_logs,
    )


@app.command("cmaes", help="Optimizes a given farm using the `CMA-ES` algorithm.")
def cmaes(
    generations: Annotated[
        int,
        typer.Argument(help="The number of total generations to perform CMA-ES."),
    ],
    seed: Annotated[
        int,
        typer.Argument(help="The seed for JAX's PRNG."),
    ],
    population: Annotated[
        int,
        typer.Option(
            "--population",
            "-p",
            help="The total population of CMA-ES. If None is given, then the "
            "population will be calculated using: 𝜆 = 4 + ⌊3 ln(n)⌋  Default is None.",
        ),
    ] = None,
    width: Annotated[
        float,
        typer.Option(
            "--width",
            "-w",
            help="The width of the farm in meters. Default is 20.5 meters.",
        ),
    ] = 20.5,
    length: Annotated[
        float,
        typer.Option(
            "--length",
            "-l",
            help="The length of the farm in meters. Default is 6.5 meters.",
        ),
    ] = 6.5,
    resolution: Annotated[
        float,
        typer.Option(
            "--resolution",
            "-r",
            help="The resolution of the image in meters. Default is 0.1 meters (10cm).",
        ),
    ] = 0.1,
    cameras: Annotated[
        int,
        typer.Option(
            "--cameras",
            "-c",
            help="The number of cameras to be placed in the farm. Default is 6.",
        ),
    ] = 6,
    camera_depth: Annotated[
        float,
        typer.Option(
            "--camera_depth",
            "-d",
            help="The depth of the camera in meters; i.e. how far can a camera should "
            "see clearly. Default is 5 meters.",
        ),
    ] = 5,
    horizontal_beam_locations: Annotated[
        list[float],
        typer.Option(
            "--horizontal_beam",
            "-hb",
            help="The horizontal beam locations that a camera can be placed on. For "
            "many locations use '-hb' multiple times. Default behaviour will create 3 "
            "beams with equal interval between them.",
        ),
    ] = None,
    vertical_beam_locations: Annotated[
        list[float],
        typer.Option(
            "--vertical_beam",
            "-vb",
            help="The vertical beam locations that a camera can be placed on. For "
            "many locations use '-vb' multiple times. Default behaviour will create 5 "
            "beams with equal interval between them.",
        ),
    ] = None,
    genotype: Annotated[
        str,
        typer.Option(
            "--genotype",
            "-gt",
            help="The genotype to be used Available genotypes: ['beams', 'xy']. "
            "Default is 'xy'.",
        ),
    ] = "xy",
) -> None:

    cmaes_optimization(
        generations,
        seed,
        population=population,
        width=width,
        length=length,
        resolution=resolution,
        cameras=cameras,
        camera_depth=camera_depth,
        horizontal_beam_locations=horizontal_beam_locations,
        vertical_beam_locations=vertical_beam_locations,
        genotype=genotype,
    )


if __name__ == "__main__":
    app()