Visualization¶

def __init__(
    self,
    output_path: str,
    dataset_name: str,
    runs: list["VizRun"],
    prepared_groundtruth: PreparedVizGroundtruth,
    observations: "ObsMetricsConfig | None" = None,
    data_root: ResolvedLocation | None = None,
):
    if not runs:
        raise ValueError("Viz requires at least one run")

    pred_dict: dict[str, dict[str, Any]] = {}
    for run in runs:
        pred_dict[run.name] = {
            "name": run.name,
            # Rollouts are written on y/x with 2-D lat/lon (see `ZarrWriter`),
            # the same layout the ground truth arrives in, so they need the
            # same rename. A no-op for runs already on lat/lon.
            "data": preserve_2d_coords(run.data),
            "ls": run.variables,
        }

    key1 = runs[0].name
    self.data_layout = prepared_groundtruth.data_layout
    levels = len(self.data_layout.depth_levels)

    var_list = {
        "vo": r"$v$ $( m/s )$",
        "uo": r"$u$ $( m/s )$",
        "thetao": r"$T$ $( ^\circ C )$",
        "tos": r"$T$ $( ^\circ C )$",
        "so": r"$so$ $( psu )$",
        "zos": r"$zos$ $( m )$",
        "KE": r"$KE$ $( J/m^2 )$",
        "OHC": r"$OHC$ $Anomaly$ $( ZJ )$",
    }

    pred_dict = process_prediction_runs(prepared_groundtruth, pred_dict)

    # Create folder paths
    self.timeseries_path = os.path.join(output_path, f"Timeseries")
    if not os.path.isdir(self.timeseries_path):
        os.makedirs(self.timeseries_path)

    self.ohc_path = os.path.join(output_path, f"OHC")
    if not os.path.isdir(self.ohc_path):
        os.makedirs(self.ohc_path)

    self.temp_path = os.path.join(output_path, f"Temperature")
    if not os.path.isdir(self.temp_path):
        os.makedirs(self.temp_path)

    self.salinity_path = os.path.join(output_path, f"Salinity")
    if not os.path.isdir(self.salinity_path):
        os.makedirs(self.salinity_path)

    self.pdfs_path = os.path.join(output_path, f"PDFs")
    if not os.path.isdir(self.pdfs_path):
        os.makedirs(self.pdfs_path)

    self.enso_path = os.path.join(output_path, f"ENSO")
    if not os.path.isdir(self.enso_path):
        os.makedirs(self.enso_path)

    self.metrics_path = os.path.join(output_path, f"Metrics")
    if not os.path.isdir(self.metrics_path):
        os.makedirs(self.metrics_path)

    self.movie_path = os.path.join(output_path, f"Movies")
    if not os.path.isdir(self.movie_path):
        os.makedirs(self.movie_path)

    clist = ["#ff807a", "#1e8685", "#ffb579", "#63c8ab"]

    # Basin masks are built on first use to avoid breaking when no correct mask
    # is available and you're not doing basin-based analyses anyway.
    self._basins: xr.Dataset = prepared_groundtruth.basins

    # Compute profile means
    with ProgressBar():
        for k in pred_dict.keys():
            logger.info("Computing profile for prediction " + k)
            pred_dict[k]["profile_prediction"] = profile_mean(
                pred_dict[k]["ds_prediction"]
            ).load()

    self.time_indices = prepared_groundtruth.time_indices
    self.data: xr.Dataset = prepared_groundtruth.data
    self.profile_groundtruth: xr.Dataset = prepared_groundtruth.profile_groundtruth
    self.pred_dict: dict[str, dict[str, Any]] = pred_dict
    self.dataset_name: str = dataset_name
    self.clist: list[str] = clist
    self.var_list: dict[str, str] = var_list
    self.levels: int = levels
    self.key1: str = key1
    self.output_path: str = output_path

    # The observation steps work from the rollouts as written, not from the
    # `pred_dict` the other steps share: they reduce through
    # `samudra.metrics`, which needs the same input `samudra.eval` gets if
    # the figures are to carry the same numbers.
    self.observations = observations
    self.obs_data_root = data_root
    self.raw_runs: list[VizRun] = list(runs)
    self.obs_path = os.path.join(output_path, "Figures_observations")
    self._obs_cache: tuple | None = None

def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)