Visualization¶

def __init__(
    self,
    output_path: str,
    dataset_name: str,
    runs: list["VizRun"],
    basins: xr.Dataset,
    groundtruth_rollout: xr.Dataset,
    time_range: slice,
):
    pred_dict: dict[str, dict[str, Any]] = {}
    for run in runs:
        pred_dict[run.name] = {
            "name": run.name,
            "data": run.data,
            "ls": run.variables,
        }

    key1 = runs[0].name
    # TODO: Support non-OM4 dataset specs in visualization.
    self.dataset_spec = build_om4_spec()
    levels = len(self.dataset_spec.depth_levels)

    groundtruth_rollout = groundtruth_rollout.sel(time=time_range)

    if "y" in groundtruth_rollout.coords:
        groundtruth_rollout = groundtruth_rollout.drop_vars(
            ["lat", "lon"], errors="ignore"
        )
        groundtruth_rollout = groundtruth_rollout.rename({"y": "lat", "x": "lon"})

    groundtruth_rollout = groundtruth_rollout.assign(
        areacello=(["lat", "lon"], spherical_area_weights(groundtruth_rollout))
    )

    # Compute real grid cell areas for physical calculations
    groundtruth_rollout["areacello_spherical"] = (
        ["lat", "lon"],
        spherical_area(groundtruth_rollout),
    )

    # This function processes the ds_groundtruth and predictions for plotting
    # The predictions are loaded into pred_dict
    data, pred_dict = process_data(
        groundtruth_rollout, pred_dict, dataset_spec=self.dataset_spec
    )

    last_index = len(data.time) - 1
    self.time_indices = [0, last_index // 2, last_index]

    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 )$",
    }

    # 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"]

    atlantic_mask0 = basins["basin_atlantic"]
    atlantic_mask = atlantic_mask0.where(atlantic_mask0["lat"] >= -32)
    atlantic_mask = process_mask(data, atlantic_mask)
    pacific_mask0 = basins["basin_pacific"]
    pacific_mask = pacific_mask0.where(
        pacific_mask0["lat"] >= -32
    )  # TODO: include this -32 masking in the basin data
    pacific_mask = process_mask(data, pacific_mask)
    indian_ocean_mask0 = basins["basin_indian"]
    indian_ocean_mask = indian_ocean_mask0.where(indian_ocean_mask0["lat"] >= -32)
    indian_ocean_mask = process_mask(data, indian_ocean_mask)
    southern_ocean_mask0 = basins["basin_southern"]
    southern_ocean_mask = process_mask(data, southern_ocean_mask0)
    arctic_mask0 = basins["basin_arctic"]
    arctic_ocean_mask = process_mask(data, arctic_mask0)

    self.basin_masks = xr.Dataset(
        {
            "Atlantic": atlantic_mask,
            "Pacific": pacific_mask,
            "Southern": southern_ocean_mask,
            "Indian": indian_ocean_mask,
            "Arctic": arctic_ocean_mask,
        }
    )

    # Compute profile means
    with ProgressBar():
        logger.info("Computing profile for ground truth " + dataset_name)
        profile_groundtruth = profile_mean(data).load()

        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.data: xr.Dataset = data
    self.profile_groundtruth: xr.Dataset = 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