roms-tools 0.20__py3-none-any.whl → 1.0.1__py3-none-any.whl

roms_tools/setup/download.py

@@ -0,0 +1,118 @@
+import pooch
+import xarray as xr
+
+# Create a Pooch object to manage the global topography data
+topo_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "etopo5.nc": "sha256:23600e422d59bbf7c3666090166a0d468c8ee16092f4f14e32c4e928fbcd627b",
+    },
+)
+
+# Create a Pooch object to manage the global SWR correction data
+correction_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "etopo5.nc": "sha256:23600e422d59bbf7c3666090166a0d468c8ee16092f4f14e32c4e928fbcd627b",
+        "SSR_correction.nc": "sha256:a170c1698e6cc2765b3f0bb51a18c6a979bc796ac3a4c014585aeede1f1f8ea0",
+    },
+)
+
+# Create a Pooch object to manage the test data
+pup_test_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-test-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "GLORYS_test_data.nc": "648f88ec29c433bcf65f257c1fb9497bd3d5d3880640186336b10ed54f7129d2",
+        "ERA5_regional_test_data.nc": "bd12ce3b562fbea2a80a3b79ba74c724294043c28dc98ae092ad816d74eac794",
+        "ERA5_global_test_data.nc": "8ed177ab64c02caf509b9fb121cf6713f286cc603b1f302f15f3f4eb0c21dc4f",
+        "TPXO_global_test_data.nc": "457bfe87a7b247ec6e04e3c7d3e741ccf223020c41593f8ae33a14f2b5255e60",
+        "TPXO_regional_test_data.nc": "11739245e2286d9c9d342dce5221e6435d2072b50028bef2e86a30287b3b4032",
+        "CESM_regional_test_data_one_time_slice.nc": "43b578ecc067c85f95d6b97ed7b9dc8da7846f07c95331c6ba7f4a3161036a17",
+        "CESM_regional_test_data_climatology.nc": "986a200029d9478fd43e6e4a8bc43e8a8f4407554893c59b5fcc2e86fd203272",
+        "CESM_surface_global_test_data_climatology.nc": "a072757110c6f7b716a98f867688ef4195a5966741d2f368201ac24617254e35",
+        "CESM_surface_global_test_data.nc": "874106ffbc8b1b220db09df1551bbb89d22439d795b4d1e5a24ee775e9a7bf6e",
+    },
+)
+
+
+def fetch_topo(topography_source: str) -> xr.Dataset:
+    """
+    Load the global topography data as an xarray Dataset.
+
+    Parameters
+    ----------
+    topography_source : str
+        The source of the topography data to be loaded. Available options:
+        - "ETOPO5"
+
+    Returns
+    -------
+    xr.Dataset
+        The global topography data as an xarray Dataset.
+    """
+    # Mapping from user-specified topography options to corresponding filenames in the registry
+    topo_dict = {"ETOPO5": "etopo5.nc"}
+
+    # Fetch the file using Pooch, downloading if necessary
+    fname = topo_data.fetch(topo_dict[topography_source])
+
+    # Load the dataset using xarray and return it
+    ds = xr.open_dataset(fname)
+    return ds
+
+
+def download_correction_data(filename: str) -> str:
+    """
+    Download the correction data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "SSR_correction.nc"
+
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+    # Fetch the file using Pooch, downloading if necessary
+    fname = correction_data.fetch(filename)
+
+    return fname
+
+
+def download_test_data(filename: str) -> str:
+    """
+    Download the test data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "GLORYS_test_data.nc"
+        - "ERA5_regional_test_data.nc"
+        - "ERA5_global_test_data.nc"
+        - "TPXO_global_test_data.nc"
+        - "TPXO_regional_test_data.nc"
+        - "CESM_regional_test_data_one_time_slice.nc"
+        - "CESM_regional_test_data_climatology.nc"
+
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+    # Fetch the file using Pooch, downloading if necessary
+    fname = pup_test_data.fetch(filename)
+
+    return fname

roms_tools/setup/initial_conditions.py

@@ -3,103 +3,80 @@ import numpy as np
 import yaml
 import importlib.metadata
 from dataclasses import dataclass, field, asdict
+from typing import Optional, Dict, Union
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.vertical_coordinate import VerticalCoordinate
 from datetime import datetime
-from roms_tools.setup.datasets import Dataset
-from roms_tools.setup.fill import fill_and_interpolate
+from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
     nan_check,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
-    extrapolate_deepest_to_bottom,
 )
+from roms_tools.setup.mixins import ROMSToolsMixins
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
 import matplotlib.pyplot as plt
 
 
 @dataclass(frozen=True, kw_only=True)
-class InitialConditions:
+class InitialConditions(ROMSToolsMixins):
     """
-    Represents initial conditions for ROMS.
+    Represents initial conditions for ROMS, including physical and biogeochemical data.
 
     Parameters
     ----------
     grid : Grid
-        Object representing the grid information.
-    vertical_coordinate: VerticalCoordinate
-        Object representing the vertical coordinate information
+        Object representing the grid information used for the model.
+    vertical_coordinate : VerticalCoordinate
+        Object representing the vertical coordinate system.
     ini_time : datetime
-        Desired initialization time.
+        The date and time at which the initial conditions are set.
+    physics_source : Dict[str, Union[str, None]]
+        Dictionary specifying the source of the physical initial condition data:
+        - "name" (str): Name of the data source (e.g., "GLORYS").
+        - "path" (str): Path to the physical data file. Can contain wildcards.
+        - "climatology" (bool): Indicates if the physical data is climatology data. Defaults to False.
+    bgc_source : Optional[Dict[str, Union[str, None]]]
+        Dictionary specifying the source of the biogeochemical (BGC) initial condition data:
+        - "name" (str): Name of the BGC data source (e.g., "CESM_REGRIDDED").
+        - "path" (str): Path to the BGC data file. Can contain wildcards.
+        - "climatology" (bool): Indicates if the BGC data is climatology data. Defaults to True.
     model_reference_date : datetime, optional
-        Reference date for the model. Default is January 1, 2000.
-    source : str, optional
-        Source of the initial condition data. Default is "GLORYS".
-    filename: str
-        Path to the source data file. Can contain wildcards.
+        The reference date for the model. Defaults to January 1, 2000.
 
     Attributes
     ----------
     ds : xr.Dataset
-        Xarray Dataset containing the initial condition data.
-
+        Xarray Dataset containing the initial condition data loaded from the specified files.
+
+    Examples
+    --------
+    >>> initial_conditions = InitialConditions(
+    ...     grid=grid,
+    ...     vertical_coordinate=vertical_coordinate,
+    ...     ini_time=datetime(2022, 1, 1),
+    ...     physics_source={"name": "GLORYS", "path": "physics_data.nc"},
+    ...     bgc_source={
+    ...         "name": "CESM_REGRIDDED",
+    ...         "path": "bgc_data.nc",
+    ...         "climatology": True,
+    ...     },
+    ... )
     """
 
     grid: Grid
     vertical_coordinate: VerticalCoordinate
     ini_time: datetime
+    physics_source: Dict[str, Union[str, None]]
+    bgc_source: Optional[Dict[str, Union[str, None]]] = None
     model_reference_date: datetime = datetime(2000, 1, 1)
-    source: str = "GLORYS"
-    filename: str
+
     ds: xr.Dataset = field(init=False, repr=False)
 
     def __post_init__(self):
 
-        # Check that the source is "GLORYS"
-        if self.source != "GLORYS":
-            raise ValueError('Only "GLORYS" is a valid option for source.')
-        if self.source == "GLORYS":
-            dims = {
-                "longitude": "longitude",
-                "latitude": "latitude",
-                "depth": "depth",
-                "time": "time",
-            }
-
-            varnames = {
-                "temp": "thetao",
-                "salt": "so",
-                "u": "uo",
-                "v": "vo",
-                "ssh": "zos",
-            }
+        self._input_checks()
+        lon, lat, angle, straddle = super().get_target_lon_lat()
 
-        data = Dataset(
-            filename=self.filename,
-            start_time=self.ini_time,
-            var_names=varnames.values(),
-            dim_names=dims,
-        )
-
-        lon = self.grid.ds.lon_rho
-        lat = self.grid.ds.lat_rho
-        angle = self.grid.ds.angle
-
-        # operate on longitudes between -180 and 180 unless ROMS domain lies at least 5 degrees in lontitude away from Greenwich meridian
-        lon = xr.where(lon > 180, lon - 360, lon)
-        straddle = True
-        if not self.grid.straddle and abs(lon).min() > 5:
-            lon = xr.where(lon < 0, lon + 360, lon)
-            straddle = False
-
-        # The following consists of two steps:
-        # Step 1: Choose subdomain of forcing data including safety margin for interpolation, and Step 2: Convert to the proper longitude range.
-        # We perform these two steps for two reasons:
-        # A) Since the horizontal dimensions consist of a single chunk, selecting a subdomain before interpolation is a lot more performant.
-        # B) Step 1 is necessary to avoid discontinuous longitudes that could be introduced by Step 2. Specifically, discontinuous longitudes
-        # can lead to artifacts in the interpolation process. Specifically, if there is a data gap if data is not global,
-        # discontinuous longitudes could result in values that appear to come from a distant location instead of producing NaNs.
-        # These NaNs are important as they can be identified and handled appropriately by the nan_check function.
+        data = self._get_data()
         data.choose_subdomain(
             latitude_range=[lat.min().values, lat.max().values],
             longitude_range=[lon.min().values, lon.max().values],
@@ -107,115 +84,134 @@ class InitialConditions:
             straddle=straddle,
         )
 
-        # interpolate onto desired grid
-        fill_dims = [dims["latitude"], dims["longitude"]]
+        vars_2d = ["zeta"]
+        vars_3d = ["temp", "salt", "u", "v"]
+        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
+        data_vars = super().process_velocities(data_vars, angle)
+
+        if self.bgc_source is not None:
+            bgc_data = self._get_bgc_data()
+            bgc_data.choose_subdomain(
+                latitude_range=[lat.min().values, lat.max().values],
+                longitude_range=[lon.min().values, lon.max().values],
+                margin=2,
+                straddle=straddle,
+            )
 
-        # 2d interpolation
-        mask = xr.where(data.ds[varnames["ssh"]].isel(time=0).isnull(), 0, 1)
-        coords = {dims["latitude"]: lat, dims["longitude"]: lon}
+            vars_2d = []
+            vars_3d = bgc_data.var_names.keys()
+            bgc_data_vars = super().regrid_data(bgc_data, vars_2d, vars_3d, lon, lat)
 
-        ssh = fill_and_interpolate(
-            data.ds[varnames["ssh"]].astype(np.float64),
-            mask,
-            fill_dims=fill_dims,
-            coords=coords,
-            method="linear",
-        )
+            # Ensure time coordinate matches if climatology is applied in one case but not the other
+            if (
+                not self.physics_source["climatology"]
+                and self.bgc_source["climatology"]
+            ):
+                for var in bgc_data_vars.keys():
+                    bgc_data_vars[var] = bgc_data_vars[var].assign_coords(
+                        {"time": data_vars["temp"]["time"]}
+                    )
 
-        # 3d interpolation
+            # Combine data variables from physical and biogeochemical sources
+            data_vars.update(bgc_data_vars)
 
-        # extrapolate deepest value all the way to bottom ("flooding")
-        for var in ["temp", "salt", "u", "v"]:
-            data.ds[varnames[var]] = extrapolate_deepest_to_bottom(
-                data.ds[varnames[var]], dims["depth"]
-            )
+        d_meta = super().get_variable_metadata()
+        ds = self._write_into_dataset(data_vars, d_meta)
 
-        mask = xr.where(data.ds[varnames["temp"]].isel(time=0).isnull(), 0, 1)
-        coords = {
-            dims["latitude"]: lat,
-            dims["longitude"]: lon,
-            dims["depth"]: self.vertical_coordinate.ds["layer_depth_rho"],
-        }
+        ds = self._add_global_metadata(ds)
 
-        # setting fillvalue_interp to None means that we allow extrapolation in the
-        # interpolation step to avoid NaNs at the surface if the lowest depth in original
-        # data is greater than zero
-        data_vars = {}
-        for var in ["temp", "salt", "u", "v"]:
-            data_vars[var] = fill_and_interpolate(
-                data.ds[varnames[var]].astype(np.float64),
-                mask,
-                fill_dims=fill_dims,
-                coords=coords,
-                method="linear",
-                fillvalue_interp=None,
-            )
+        ds["zeta"].load()
+        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
 
-        # rotate to grid orientation
-        u_rot = data_vars["u"] * np.cos(angle) + data_vars["v"] * np.sin(angle)
-        v_rot = data_vars["v"] * np.cos(angle) - data_vars["u"] * np.sin(angle)
+        object.__setattr__(self, "ds", ds)
 
-        # interpolate to u- and v-points
-        u = interpolate_from_rho_to_u(u_rot)
-        v = interpolate_from_rho_to_v(v_rot)
+    def _input_checks(self):
 
-        # 3d masks for ROMS domain
-        umask = self.grid.ds.mask_u.expand_dims({"s_rho": u.s_rho})
-        vmask = self.grid.ds.mask_v.expand_dims({"s_rho": v.s_rho})
+        if "name" not in self.physics_source.keys():
+            raise ValueError("`physics_source` must include a 'name'.")
+        if "path" not in self.physics_source.keys():
+            raise ValueError("`physics_source` must include a 'path'.")
+        # set self.physics_source["climatology"] to False if not provided
+        object.__setattr__(
+            self,
+            "physics_source",
+            {
+                **self.physics_source,
+                "climatology": self.physics_source.get("climatology", False),
+            },
+        )
+        if self.bgc_source is not None:
+            if "name" not in self.bgc_source.keys():
+                raise ValueError(
+                    "`bgc_source` must include a 'name' if it is provided."
+                )
+            if "path" not in self.bgc_source.keys():
+                raise ValueError(
+                    "`bgc_source` must include a 'path' if it is provided."
+                )
+            # set self.bgc_source["climatology"] to True if not provided
+            object.__setattr__(
+                self,
+                "bgc_source",
+                {
+                    **self.bgc_source,
+                    "climatology": self.bgc_source.get("climatology", True),
+                },
+            )
 
-        u = u * umask
-        v = v * vmask
+    def _get_data(self):
 
-        # Compute barotropic velocity
-        # thicknesses
-        dz = -self.vertical_coordinate.ds["interface_depth_rho"].diff(dim="s_w")
-        dz = dz.rename({"s_w": "s_rho"})
-        # thicknesses at u- and v-points
-        dzu = interpolate_from_rho_to_u(dz)
-        dzv = interpolate_from_rho_to_v(dz)
+        if self.physics_source["name"] == "GLORYS":
+            data = GLORYSDataset(
+                filename=self.physics_source["path"],
+                start_time=self.ini_time,
+                climatology=self.physics_source["climatology"],
+            )
+        else:
+            raise ValueError(
+                'Only "GLORYS" is a valid option for physics_source["name"].'
+            )
+        return data
 
-        ubar = (dzu * u).sum(dim="s_rho") / dzu.sum(dim="s_rho")
-        vbar = (dzv * v).sum(dim="s_rho") / dzv.sum(dim="s_rho")
+    def _get_bgc_data(self):
 
-        # save in new dataset
-        ds = xr.Dataset()
+        if self.bgc_source["name"] == "CESM_REGRIDDED":
 
-        ds["temp"] = data_vars["temp"].astype(np.float32)
-        ds["temp"].attrs["long_name"] = "Potential temperature"
-        ds["temp"].attrs["units"] = "Celsius"
+            bgc_data = CESMBGCDataset(
+                filename=self.bgc_source["path"],
+                start_time=self.ini_time,
+                climatology=self.bgc_source["climatology"],
+            )
+            bgc_data.post_process()
+        else:
+            raise ValueError(
+                'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
+            )
 
-        ds["salt"] = data_vars["salt"].astype(np.float32)
-        ds["salt"].attrs["long_name"] = "Salinity"
-        ds["salt"].attrs["units"] = "PSU"
+        return bgc_data
 
-        ds["zeta"] = ssh.astype(np.float32)
-        ds["zeta"].attrs["long_name"] = "Free surface"
-        ds["zeta"].attrs["units"] = "m"
+    def _write_into_dataset(self, data_vars, d_meta):
 
-        ds["u"] = u.astype(np.float32)
-        ds["u"].attrs["long_name"] = "u-flux component"
-        ds["u"].attrs["units"] = "m/s"
+        # save in new dataset
+        ds = xr.Dataset()
 
-        ds["v"] = v.astype(np.float32)
-        ds["v"].attrs["long_name"] = "v-flux component"
-        ds["v"].attrs["units"] = "m/s"
+        for var in data_vars.keys():
+            ds[var] = data_vars[var].astype(np.float32)
+            ds[var].attrs["long_name"] = d_meta[var]["long_name"]
+            ds[var].attrs["units"] = d_meta[var]["units"]
 
         # initialize vertical velocity to zero
         ds["w"] = xr.zeros_like(
             self.vertical_coordinate.ds["interface_depth_rho"].expand_dims(
-                time=ds[dims["time"]]
+                time=data_vars["u"].time
             )
         ).astype(np.float32)
-        ds["w"].attrs["long_name"] = "w-flux component"
-        ds["w"].attrs["units"] = "m/s"
+        ds["w"].attrs["long_name"] = d_meta["w"]["long_name"]
+        ds["w"].attrs["units"] = d_meta["w"]["units"]
 
-        ds["ubar"] = ubar.transpose(dims["time"], "eta_rho", "xi_u").astype(np.float32)
-        ds["ubar"].attrs["long_name"] = "vertically integrated u-flux component"
-        ds["ubar"].attrs["units"] = "m/s"
+        return ds
 
-        ds["vbar"] = vbar.transpose(dims["time"], "eta_v", "xi_rho").astype(np.float32)
-        ds["vbar"].attrs["long_name"] = "vertically integrated v-flux component"
-        ds["vbar"].attrs["units"] = "m/s"
+    def _add_global_metadata(self, ds):
 
         ds = ds.assign_coords(
             {
@@ -235,10 +231,9 @@ class InitialConditions:
         ds.attrs["roms_tools_version"] = roms_tools_version
         ds.attrs["ini_time"] = str(self.ini_time)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
-        ds.attrs["source"] = self.source
-
-        if dims["time"] != "time":
-            ds = ds.rename({dims["time"]: "time"})
+        ds.attrs["physical_source"] = self.physics_source["name"]
+        if self.bgc_source is not None:
+            ds.attrs["bgc_source"] = self.bgc_source["name"]
 
         # Translate the time coordinate to days since the model reference date
         model_reference_date = np.datetime64(self.model_reference_date)
@@ -248,22 +243,19 @@ class InitialConditions:
         ds = ds.assign_coords(ocean_time=("time", np.float32(ocean_time)))
         ds["ocean_time"].attrs[
             "long_name"
-        ] = f"time since {np.datetime_as_string(model_reference_date, unit='D')}"
+        ] = f"seconds since {np.datetime_as_string(model_reference_date, unit='s')}"
         ds["ocean_time"].attrs["units"] = "seconds"
 
-        ds["theta_s"] = self.vertical_coordinate.ds["theta_s"]
-        ds["theta_b"] = self.vertical_coordinate.ds["theta_b"]
-        ds["Tcline"] = self.vertical_coordinate.ds["Tcline"]
-        ds["hc"] = self.vertical_coordinate.ds["hc"]
+        ds.attrs["theta_s"] = self.vertical_coordinate.ds["theta_s"].item()
+        ds.attrs["theta_b"] = self.vertical_coordinate.ds["theta_b"].item()
+        ds.attrs["Tcline"] = self.vertical_coordinate.ds["Tcline"].item()
+        ds.attrs["hc"] = self.vertical_coordinate.ds["hc"].item()
         ds["sc_r"] = self.vertical_coordinate.ds["sc_r"]
         ds["Cs_r"] = self.vertical_coordinate.ds["Cs_r"]
 
         ds = ds.drop_vars(["s_rho"])
 
-        object.__setattr__(self, "ds", ds)
-
-        ds["zeta"].load()
-        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+        return ds
 
     def plot(
         self,
@@ -289,6 +281,38 @@ class InitialConditions:
             - "w": w-flux component.
             - "ubar": Vertically integrated u-flux component.
             - "vbar": Vertically integrated v-flux component.
+            - "PO4": Dissolved Inorganic Phosphate (mmol/m³).
+            - "NO3": Dissolved Inorganic Nitrate (mmol/m³).
+            - "SiO3": Dissolved Inorganic Silicate (mmol/m³).
+            - "NH4": Dissolved Ammonia (mmol/m³).
+            - "Fe": Dissolved Inorganic Iron (mmol/m³).
+            - "Lig": Iron Binding Ligand (mmol/m³).
+            - "O2": Dissolved Oxygen (mmol/m³).
+            - "DIC": Dissolved Inorganic Carbon (mmol/m³).
+            - "DIC_ALT_CO2": Dissolved Inorganic Carbon, Alternative CO2 (mmol/m³).
+            - "ALK": Alkalinity (meq/m³).
+            - "ALK_ALT_CO2": Alkalinity, Alternative CO2 (meq/m³).
+            - "DOC": Dissolved Organic Carbon (mmol/m³).
+            - "DON": Dissolved Organic Nitrogen (mmol/m³).
+            - "DOP": Dissolved Organic Phosphorus (mmol/m³).
+            - "DOPr": Refractory Dissolved Organic Phosphorus (mmol/m³).
+            - "DONr": Refractory Dissolved Organic Nitrogen (mmol/m³).
+            - "DOCr": Refractory Dissolved Organic Carbon (mmol/m³).
+            - "zooC": Zooplankton Carbon (mmol/m³).
+            - "spChl": Small Phytoplankton Chlorophyll (mg/m³).
+            - "spC": Small Phytoplankton Carbon (mmol/m³).
+            - "spP": Small Phytoplankton Phosphorous (mmol/m³).
+            - "spFe": Small Phytoplankton Iron (mmol/m³).
+            - "spCaCO3": Small Phytoplankton CaCO3 (mmol/m³).
+            - "diatChl": Diatom Chlorophyll (mg/m³).
+            - "diatC": Diatom Carbon (mmol/m³).
+            - "diatP": Diatom Phosphorus (mmol/m³).
+            - "diatFe": Diatom Iron (mmol/m³).
+            - "diatSi": Diatom Silicate (mmol/m³).
+            - "diazChl": Diazotroph Chlorophyll (mg/m³).
+            - "diazC": Diazotroph Carbon (mmol/m³).
+            - "diazP": Diazotroph Phosphorus (mmol/m³).
+            - "diazFe": Diazotroph Iron (mmol/m³).
         s : int, optional
             The index of the vertical layer to plot. Default is None.
         eta : int, optional
@@ -306,9 +330,9 @@ class InitialConditions:
         Raises
         ------
         ValueError
-            If the specified varname is not one of the valid options.
-            If field is 3D and none of s_rho, eta, xi are specified.
-            If field is 2D and both eta and xi are specified.
+            If the specified `varname` is not one of the valid options.
+            If the field specified by `varname` is 3D and none of `s`, `eta`, or `xi` are specified.
+            If the field specified by `varname` is 2D and both `eta` and `xi` are specified.
         """
 
         if len(self.ds[varname].squeeze().dims) == 3 and not any(
@@ -376,7 +400,10 @@ class InitialConditions:
         else:
             vmax = field.max().values
             vmin = field.min().values
-            cmap = plt.colormaps.get_cmap("YlOrRd")
+            if varname in ["temp", "salt"]:
+                cmap = plt.colormaps.get_cmap("YlOrRd")
+            else:
+                cmap = plt.colormaps.get_cmap("YlGn")
         cmap.set_bad(color="gray")
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
@@ -454,12 +481,14 @@ class InitialConditions:
 
         initial_conditions_data = {
             "InitialConditions": {
-                "filename": self.filename,
+                "physics_source": self.physics_source,
                 "ini_time": self.ini_time.isoformat(),
                 "model_reference_date": self.model_reference_date.isoformat(),
-                "source": self.source,
             }
         }
+        # Include bgc_source if it's not None
+        if self.bgc_source is not None:
+            initial_conditions_data["InitialConditions"]["bgc_source"] = self.bgc_source
 
         yaml_data = {
             **grid_yaml_data,