roms-tools 0.0.2__py3-none-any.whl → 0.1.0__py3-none-any.whl

roms_tools/setup/grid.py

@@ -1,51 +1,93 @@
 import copy
 from dataclasses import dataclass, field
-from datetime import date
 
 import numpy as np
 import xarray as xr
 
-from typing import Any
 
+from roms_tools.setup.topography import _add_topography_and_mask
+from roms_tools.setup.plot import _plot
 
 RADIUS_OF_EARTH = 6371315.0  # in m
 
 
-# TODO lat_rho and lon_rho should be coordinate variables
-
 # TODO should we store an xgcm.Grid object instead of an xarray Dataset? Or even subclass xgcm.Grid?
 
 
 @dataclass(frozen=True, kw_only=True)
 class Grid:
     """
-    A single ROMS grid.
-
-    Used for creating, plotting, and then saving a new ROMS domain grid.
-
-    Parameters
-    ----------
-    nx
-        Number of grid points in the x-direction
-    ny
-        Number of grid points in the y-direction
-    size_x
-        Domain size in the x-direction (in km?)
-    size_y
-        Domain size in the y-direction (in km?)
-    center_lon
-        Longitude of grid center
-    center_lat
-        Latitude of grid center
-    rot
-        Rotation of grid x-direction from lines of constant latitude.
-        Measured in degrees, with positive values meaning a counterclockwise rotation.
-        The default is 0, which means that the x-direction of the grid x-direction is aligned with lines of constant latitude.
-
-    Raises
-    ------
-    ValueError
-        If you try to create a grid which crosses the Greenwich Meridian
+     A single ROMS grid.
+
+     Used for creating, plotting, and then saving a new ROMS domain grid.
+
+     Parameters
+     ----------
+     nx : int
+         Number of grid points in the x-direction.
+     ny : int
+         Number of grid points in the y-direction.
+     size_x : float
+         Domain size in the x-direction (in kilometers).
+     size_y : float
+         Domain size in the y-direction (in kilometers).
+     center_lon : float
+         Longitude of grid center.
+     center_lat : float
+         Latitude of grid center.
+     rot : float, optional
+         Rotation of grid x-direction from lines of constant latitude, measured in degrees.
+         Positive values represent a counterclockwise rotation.
+         The default is 0, which means that the x-direction of the grid is aligned with lines of constant latitude.
+    topography_source : str, optional
+        Specifies the data source to use for the topography. Options are
+        "etopo5". The default is "etopo5".
+    smooth_factor : float, optional
+        The smoothing factor used in the domain-wide Gaussian smoothing of the
+        topography. Smaller values result in less smoothing, while larger
+        values produce more smoothing. The default is 8.
+    hmin : float, optional
+        The minimum ocean depth (in meters). The default is 5.
+    rmax : float, optional
+        The maximum slope parameter (in meters). This parameter controls
+        the local smoothing of the topography. Smaller values result in
+        smoother topography, while larger values preserve more detail.
+        The default is 0.2.
+
+     Attributes
+     ----------
+     nx : int
+         Number of grid points in the x-direction.
+     ny : int
+         Number of grid points in the y-direction.
+     size_x : float
+         Domain size in the x-direction (in kilometers).
+     size_y : float
+         Domain size in the y-direction (in kilometers).
+     center_lon : float
+         Longitude of grid center.
+     center_lat : float
+         Latitude of grid center.
+     rot : float
+         Rotation of grid x-direction from lines of constant latitude.
+     topography_source : str
+         Data source used for the topography.
+     smooth_factor : int
+         Smoothing factor used in the domain-wide Gaussian smoothing of the topography.
+     hmin : float
+         Minimum ocean depth (in meters).
+     rmax : float
+         Maximum slope parameter (in meters).
+     ds : xr.Dataset
+         The xarray Dataset containing the grid data.
+     straddle : bool
+         Indicates if the Greenwich meridian (0° longitude) intersects the domain.
+         `True` if it does, `False` otherwise.
+
+     Raises
+     ------
+     ValueError
+         If you try to create a grid with domain size larger than 20000 km.
     """
 
     nx: int
@@ -55,7 +97,12 @@ class Grid:
     center_lon: float
     center_lat: float
     rot: float = 0
+    topography_source: str = "etopo5"
+    smooth_factor: int = 8
+    hmin: float = 5.0
+    rmax: float = 0.2
     ds: xr.Dataset = field(init=False, repr=False)
+    straddle: bool = field(init=False, repr=False)
 
     def __post_init__(self):
         ds = _make_grid_ds(
@@ -71,6 +118,82 @@ class Grid:
         # see https://stackoverflow.com/questions/53756788/how-to-set-the-value-of-dataclass-field-in-post-init-when-frozen-true
         object.__setattr__(self, "ds", ds)
 
+        # Update self.ds with topography and mask information
+        self.add_topography_and_mask(
+            topography_source=self.topography_source,
+            smooth_factor=self.smooth_factor,
+            hmin=self.hmin,
+            rmax=self.rmax,
+        )
+
+        # Check if the Greenwich meridian goes through the domain.
+        self._straddle()
+
+    def add_topography_and_mask(
+        self, topography_source="etopo5", smooth_factor=8, hmin=5.0, rmax=0.2
+    ) -> None:
+        """
+        Add topography and mask to the grid dataset.
+
+        This method processes the topography data and generates a land/sea mask.
+        It applies several steps, including interpolating topography, smoothing
+        the topography over the entire domain and locally, and filling in enclosed basins. The
+        processed topography and mask are added to the grid's dataset as new variables.
+
+        Parameters
+        ----------
+        topography_source : str, optional
+            Specifies the data source to use for the topography. Options are
+            "etopo5". The default is "etopo5".
+        smooth_factor : float, optional
+            The smoothing factor used in the domain-wide Gaussian smoothing of the
+            topography. Smaller values result in less smoothing, while larger
+            values produce more smoothing. The default is 8.
+        hmin : float, optional
+            The minimum ocean depth (in meters). The default is 5.
+        rmax : float, optional
+            The maximum slope parameter (in meters). This parameter controls
+            the local smoothing of the topography. Smaller values result in
+            smoother topography, while larger values preserve more detail.
+            The default is 0.2.
+
+        Returns
+        -------
+        None
+            This method modifies the dataset in place and does not return a value.
+        """
+
+        ds = _add_topography_and_mask(
+            self.ds, topography_source, smooth_factor, hmin, rmax
+        )
+        # Assign the updated dataset back to the frozen dataclass
+        object.__setattr__(self, "ds", ds)
+
+    def compute_bathymetry_laplacian(self):
+        """
+        Compute the Laplacian of the 'h' field in the provided grid dataset.
+
+        Adds:
+        xarray.DataArray: The Laplacian of the 'h' field as a new variable in the dataset self.ds.
+        """
+
+        # Extract the 'h' field and grid spacing variables
+        h = self.ds.h
+        pm = self.ds.pm  # Reciprocal of grid spacing in x-direction
+        pn = self.ds.pn  # Reciprocal of grid spacing in y-direction
+
+        # Compute second derivatives using finite differences
+        d2h_dx2 = (h.shift(xi_rho=-1) - 2 * h + h.shift(xi_rho=1)) * pm**2
+        d2h_dy2 = (h.shift(eta_rho=-1) - 2 * h + h.shift(eta_rho=1)) * pn**2
+
+        # Compute the Laplacian by summing second derivatives
+        laplacian_h = d2h_dx2 + d2h_dy2
+
+        # Add the Laplacian as a new variable in the dataset
+        self.ds["h_laplacian"] = laplacian_h
+        self.ds["h_laplacian"].attrs["long_name"] = "Laplacian of final bathymetry"
+        self.ds["h_laplacian"].attrs["units"] = "1/m"
+
     def save(self, filepath: str) -> None:
         """
         Save the grid information to a netCDF4 file.
@@ -81,78 +204,158 @@ class Grid:
         """
         self.ds.to_netcdf(filepath)
 
-    def from_file(self, filepath: str) -> "Grid":
+    @classmethod
+    def from_file(cls, filepath: str) -> "Grid":
         """
-        Open an existing grid from a file.
+        Create a Grid instance from an existing file.
 
         Parameters
         ----------
-        filepath
+        filepath : str
+            Path to the file containing the grid information.
+
+        Returns
+        -------
+        Grid
+            A new instance of Grid populated with data from the file.
         """
-        # TODO set other parameters that were saved into the file, because every parameter we need gets saved.
-        # TODO actually we will need to deduce size_x and size_y from the file, that's annoying.
-        self.ds = xr.open_dataset(filepath)
-        raise NotImplementedError()
+        # Load the dataset from the file
+        ds = xr.open_dataset(filepath)
+
+        # Create a new Grid instance without calling __init__ and __post_init__
+        grid = cls.__new__(cls)
+
+        # Set the dataset for the grid instance
+        object.__setattr__(grid, "ds", ds)
+
+        # Check if the Greenwich meridian goes through the domain.
+        grid._straddle()
+
+        # Manually set the remaining attributes by extracting parameters from dataset
+        object.__setattr__(grid, "nx", ds.sizes["xi_rho"] - 2)
+        object.__setattr__(grid, "ny", ds.sizes["eta_rho"] - 2)
+        object.__setattr__(grid, "center_lon", ds["tra_lon"].values.item())
+        object.__setattr__(grid, "center_lat", ds["tra_lat"].values.item())
+        object.__setattr__(grid, "rot", ds["rotate"].values.item())
+
+        for attr in [
+            "size_x",
+            "size_y",
+            "topography_source",
+            "smooth_factor",
+            "hmin",
+            "rmax",
+        ]:
+            if attr in ds.attrs:
+                object.__setattr__(grid, attr, ds.attrs[attr])
+
+        return grid
+
+    # override __repr__ method to only print attributes that are actually set
+    def __repr__(self) -> str:
+        cls = self.__class__
+        cls_name = cls.__name__
+        # Create a dictionary of attribute names and values, filtering out those that are not set and 'ds'
+        attr_dict = {
+            k: v for k, v in self.__dict__.items() if k != "ds" and v is not None
+        }
+        attr_str = ", ".join(f"{k}={v!r}" for k, v in attr_dict.items())
+        return f"{cls_name}({attr_str})"
+
+    # def to_xgcm() -> Any:
+    #    # TODO we could convert the dataset to an xgcm.Grid object and return here?
+    #    raise NotImplementedError()
+
+    def _straddle(self) -> None:
+        """
+        Check if the Greenwich meridian goes through the domain.
+
+        This method sets the `straddle` attribute to `True` if the Greenwich meridian
+        (0° longitude) intersects the domain defined by `lon_rho`. Otherwise, it sets
+        the `straddle` attribute to `False`.
 
-    def to_xgcm() -> Any:
-        # TODO we could convert the dataset to an xgcm.Grid object and return here?
-        raise NotImplementedError()
+        The check is based on whether the longitudinal differences between adjacent
+        points exceed 300 degrees, indicating a potential wraparound of longitude.
+        """
+
+        if (
+            np.abs(self.ds.lon_rho.diff("xi_rho")).max() > 300
+            or np.abs(self.ds.lon_rho.diff("eta_rho")).max() > 300
+        ):
+            object.__setattr__(self, "straddle", True)
+        else:
+            object.__setattr__(self, "straddle", False)
 
     def plot(self, bathymetry: bool = False) -> None:
         """
         Plot the grid.
 
-        Requires cartopy and matplotlib.
-
         Parameters
         ----------
-        bathymetry: bool
+        bathymetry : bool
             Whether or not to plot the bathymetry. Default is False.
+
+        Returns
+        -------
+        None
+            This method does not return any value. It generates and displays a plot.
+
         """
 
-        # TODO optionally plot topography on top?
         if bathymetry:
-            raise NotImplementedError()
+            kwargs = {"cmap": "YlGnBu"}
+            _plot(
+                self.ds,
+                field=self.ds.h.where(self.ds.mask_rho),
+                straddle=self.straddle,
+                kwargs=kwargs,
+            )
+        else:
+            _plot(self.ds, straddle=self.straddle)
+
+    def coarsen(self):
+        """
+        Update the grid by adding grid variables that are coarsened versions of the original
+        fine-resoluion grid variables. The coarsening is by a factor of two.
+
+        The specific variables being coarsened are:
+        - `lon_rho` -> `lon_coarse`: Longitude at rho points.
+        - `lat_rho` -> `lat_coarse`: Latitude at rho points.
+        - `h` -> `h_coarse`: Bathymetry (depth).
+        - `angle` -> `angle_coarse`: Angle between the xi axis and true east.
+        - `mask_rho` -> `mask_coarse`: Land/sea mask at rho points.
+
+        Returns
+        -------
+        None
+
+        Modifies
+        --------
+        self.ds : xr.Dataset
+            The dataset attribute of the Grid instance is updated with the new coarser variables.
+        """
+        d = {
+            "lon_rho": "lon_coarse",
+            "lat_rho": "lat_coarse",
+            "h": "h_coarse",
+            "angle": "angle_coarse",
+            "mask_rho": "mask_coarse",
+        }
 
-        import cartopy.crs as ccrs
-        import matplotlib.pyplot as plt
+        for fine_var, coarse_var in d.items():
+            fine_field = self.ds[fine_var]
+            if self.straddle and fine_var == "lon_rho":
+                fine_field = xr.where(fine_field > 180, fine_field - 360, fine_field)
 
-        lon_deg = (self.ds["lon_rho"] - 360).values
-        lat_deg = self.ds["lat_rho"].values
+            coarse_field = _f2c(fine_field)
+            if fine_var == "lon_rho":
+                coarse_field = xr.where(
+                    coarse_field < 0, coarse_field + 360, coarse_field
+                )
 
-        # Define projections
-        geodetic = ccrs.Geodetic()
-        trans = ccrs.NearsidePerspective(
-            central_longitude=np.mean(lon_deg), central_latitude=np.mean(lat_deg)
-        )
+            self.ds[coarse_var] = coarse_field
 
-        # find corners
-        (lo1, la1) = (lon_deg[0, 0], lat_deg[0, 0])
-        (lo2, la2) = (lon_deg[0, -1], lat_deg[0, -1])
-        (lo3, la3) = (lon_deg[-1, -1], lat_deg[-1, -1])
-        (lo4, la4) = (lon_deg[-1, 0], lat_deg[-1, 0])
-
-        # transform coordinates to projected space
-        lo1t, la1t = trans.transform_point(lo1, la1, geodetic)
-        lo2t, la2t = trans.transform_point(lo2, la2, geodetic)
-        lo3t, la3t = trans.transform_point(lo3, la3, geodetic)
-        lo4t, la4t = trans.transform_point(lo4, la4, geodetic)
-
-        plt.figure(figsize=(10, 10))
-        ax = plt.axes(projection=trans)
-
-        ax.plot(
-            [lo1t, lo2t, lo3t, lo4t, lo1t],
-            [la1t, la2t, la3t, la4t, la1t],
-            "ro-",
-        )
-
-        ax.coastlines(
-            resolution="50m", linewidth=0.5, color="black"
-        )  # add map of coastlines
-        ax.gridlines()
-
-        plt.show()
+        self.ds["mask_coarse"] = xr.where(self.ds["mask_coarse"] > 0.5, 1, 0)
 
 
 def _make_grid_ds(
@@ -164,195 +367,145 @@ def _make_grid_ds(
     center_lat: float,
     rot: float,
 ) -> xr.Dataset:
+    _raise_if_domain_size_too_large(size_x, size_y)
 
-    if size_y > size_x:
-        domain_length, domain_width = size_x * 1e3, size_y * 1e3  # in m
-        nl, nw = nx, ny
-    else:
-        domain_length, domain_width = size_y * 1e3, size_x * 1e3  # in m
-        nl, nw = ny, nx
-
-    initial_lon_lat_vars = _make_initial_lon_lat_ds(domain_length, domain_width, nl, nw)
+    initial_lon_lat_vars = _make_initial_lon_lat_ds(size_x, size_y, nx, ny)
 
+    # rotate coordinate system
     rotated_lon_lat_vars = _rotate(*initial_lon_lat_vars, rot)
+    lon, *_ = rotated_lon_lat_vars
 
-    lon2, *_ = rotated_lon_lat_vars
-
-    _raise_if_crosses_greenwich_meridian(lon2, center_lon)
-
+    # translate coordinate system
     translated_lon_lat_vars = _translate(*rotated_lon_lat_vars, center_lat, center_lon)
+    lon, lat, lonu, latu, lonv, latv, lonq, latq = translated_lon_lat_vars
 
-    lon4, lat4, lonu, latu, lonv, latv, lone, late = translated_lon_lat_vars
-
-    pm, pn = _compute_coordinate_metrics(lon4, lonu, latu, lonv, latv)
-
-    ang = _compute_angle(lon4, lonu, latu, lone)
-
-    ds = _create_grid_ds(
-        nx,
-        ny,
-        lon4,
-        lat4,
-        pm,
-        pn,
-        ang,
-        size_x,
-        size_y,
-        rot,
-        center_lon,
-        center_lat,
-        lone,
-        late,
-    )
-
-    # TODO topography
-    # ds = _make_topography(ds)
+    # compute 1/dx and 1/dy
+    pm, pn = _compute_coordinate_metrics(lon, lonu, latu, lonv, latv)
 
-    ds = _add_global_metadata(ds, nx, ny, size_x, size_y, center_lon, center_lat, rot)
+    # compute angle of local grid positive x-axis relative to east
+    ang = _compute_angle(lon, lonu, latu, lonq)
 
-    return ds
+    ds = _create_grid_ds(lon, lat, pm, pn, ang, rot, center_lon, center_lat)
 
+    ds = _add_global_metadata(ds, size_x, size_y)
 
-def _raise_if_crosses_greenwich_meridian(lon, center_lon):
-    # We have to do this before the grid is translated because we don't trust the grid creation routines in that case.
+    return ds
 
-    # TODO it would be nice to handle this case, but we first need to know what ROMS expects / can handle.
 
-    # TODO what about grids which cross the international dateline?
+def _raise_if_domain_size_too_large(size_x, size_y):
+    threshold = 20000
+    if size_x > threshold or size_y > threshold:
+        raise ValueError("Domain size has to be smaller than %g km" % threshold)
 
-    if np.min(lon + center_lon) < 0 < np.max(lon + center_lon):
-        raise ValueError("Grid cannot cross Greenwich Meridian")
 
+def _make_initial_lon_lat_ds(size_x, size_y, nx, ny):
+    # Mercator projection around the equator
 
-def _make_initial_lon_lat_ds(domain_length, domain_width, nl, nw):
+    # initially define the domain to be longer in x-direction (dimension "length")
+    # than in y-direction (dimension "width") to keep grid distortion minimal
+    if size_y > size_x:
+        domain_length, domain_width = size_y * 1e3, size_x * 1e3  # in m
+        nl, nw = ny, nx
+    else:
+        domain_length, domain_width = size_x * 1e3, size_y * 1e3  # in m
+        nl, nw = nx, ny
 
-    domain_length_in_degrees_longitude = domain_length / RADIUS_OF_EARTH
-    domain_width_in_degrees_latitude = domain_width / RADIUS_OF_EARTH
+    domain_length_in_degrees = domain_length / RADIUS_OF_EARTH
+    domain_width_in_degrees = domain_width / RADIUS_OF_EARTH
 
-    longitude_array_1d_in_degrees = (
-        domain_length_in_degrees_longitude * np.arange(-0.5, nl + 1.5, 1) / nl
-        - domain_length_in_degrees_longitude / 2
+    # 1d array describing the longitudes at cell centers
+    x = np.arange(-0.5, nl + 1.5, 1)
+    lon_array_1d_in_degrees = (
+        domain_length_in_degrees * x / nl - domain_length_in_degrees / 2
+    )
+    # 1d array describing the longitudes at cell corners (or vorticity points "q")
+    xq = np.arange(-1, nl + 2, 1)
+    lonq_array_1d_in_degrees_q = (
+        domain_length_in_degrees * xq / nl - domain_length_in_degrees / 2
     )
 
-    # TODO I don't fully understand what this piece of code achieves
-    mul = 1.0
-    for it in range(1, 101):
+    # convert degrees latitude to y-coordinate using Mercator projection
+    y1 = np.log(np.tan(np.pi / 4 - domain_width_in_degrees / 4))
+    y2 = np.log(np.tan(np.pi / 4 + domain_width_in_degrees / 4))
 
-        # convert degrees latitude to y-coordinate using Mercator projection
-        y1 = np.log(np.tan(np.pi / 4 - domain_width_in_degrees_latitude / 4))
-        y2 = np.log(np.tan(np.pi / 4 + domain_width_in_degrees_latitude / 4))
+    # linearly space points in y-space
+    y = (y2 - y1) * np.arange(-0.5, nw + 1.5, 1) / nw + y1
+    yq = (y2 - y1) * np.arange(-1, nw + 2) / nw + y1
 
-        # linearly space points in y-space
-        y = (y2 - y1) * np.arange(-0.5, nw + 1.5, 1) / nw + y1
+    # inverse Mercator projections
+    lat_array_1d_in_degrees = np.arctan(np.sinh(y))
+    latq_array_1d_in_degrees = np.arctan(np.sinh(yq))
 
-        # convert back to longitude using inverse Mercator projection
-        # lat1d = 2*np.arctan(np.exp(y)) - np.pi/2
-        latitude_array_1d_in_degrees = np.arctan(np.sinh(y))
+    # 2d grid at cell centers
+    lon, lat = np.meshgrid(lon_array_1d_in_degrees, lat_array_1d_in_degrees)
+    # 2d grid at cell corners
+    lonq, latq = np.meshgrid(lonq_array_1d_in_degrees_q, latq_array_1d_in_degrees)
 
-        # find width and height of new grid at central grid point in degrees
-        latitude_array_1d_in_degrees_cen = 0.5 * (
-            latitude_array_1d_in_degrees[int(np.round(nw / 2) + 1)]
-            - latitude_array_1d_in_degrees[int(np.round(nw / 2) - 1)]
-        )
-        longitude_array_1d_in_degrees_cen = domain_length_in_degrees_longitude / nl
-
-        # scale the domain width in degreees latitude somehow?
-        mul = (
-            latitude_array_1d_in_degrees_cen
-            / longitude_array_1d_in_degrees_cen
-            * domain_length_in_degrees_longitude
-            / domain_width_in_degrees_latitude
-            * nw
-            / nl
-        )
-        latitude_array_1d_in_degrees = latitude_array_1d_in_degrees / mul
+    if size_y > size_x:
+        # Rotate grid by 90 degrees because until here the grid has been defined
+        # to be longer in x-direction than in y-direction
 
-    # TODO what does the 'e' suffix mean?
-    lon1de = (
-        domain_length_in_degrees_longitude * np.arange(-1, nl + 2, 1) / nl
-        - domain_length_in_degrees_longitude / 2
-    )
-    ye = (y2 - y1) * np.arange(-1, nw + 2) / nw + y1
-    # lat1de = 2 * np.arctan(np.exp(ye)) - np.pi/2
-    lat1de = np.arctan(np.sinh(ye))
-    lat1de = lat1de / mul
+        lon, lat = _rot_sphere(lon, lat, 90)
+        lonq, latq = _rot_sphere(lonq, latq, 90)
 
-    lon1, lat1 = np.meshgrid(
-        longitude_array_1d_in_degrees, latitude_array_1d_in_degrees
-    )
-    lone, late = np.meshgrid(lon1de, lat1de)
-    lonu = 0.5 * (lon1[:, :-1] + lon1[:, 1:])
-    latu = 0.5 * (lat1[:, :-1] + lat1[:, 1:])
-    lonv = 0.5 * (lon1[:-1, :] + lon1[1:, :])
-    latv = 0.5 * (lat1[:-1, :] + lat1[1:, :])
-
-    if domain_length > domain_width:
-        # Rotate grid 90 degrees so that the width is now longer than the length
-
-        lon1, lat1 = rot_sphere(lon1, lat1, 90)
-        lonu, latu = rot_sphere(lonu, latu, 90)
-        lonv, latv = rot_sphere(lonv, latv, 90)
-        lone, late = rot_sphere(lone, late, 90)
-
-        lon1 = np.transpose(np.flip(lon1, 0))
-        lat1 = np.transpose(np.flip(lat1, 0))
-        lone = np.transpose(np.flip(lone, 0))
-        late = np.transpose(np.flip(late, 1))
-
-        lonu_tmp = np.transpose(np.flip(lonv, 0))
-        latu_tmp = np.transpose(np.flip(latv, 0))
-        lonv = np.transpose(np.flip(lonu, 0))
-        latv = np.transpose(np.flip(latu, 0))
-        lonu = lonu_tmp
-        latu = latu_tmp
+        lon = np.transpose(np.flip(lon, 0))
+        lat = np.transpose(np.flip(lat, 0))
+        lonq = np.transpose(np.flip(lonq, 0))
+        latq = np.transpose(np.flip(latq, 0))
+
+    # infer longitudes and latitudes at u- and v-points
+    lonu = 0.5 * (lon[:, :-1] + lon[:, 1:])
+    latu = 0.5 * (lat[:, :-1] + lat[:, 1:])
+    lonv = 0.5 * (lon[:-1, :] + lon[1:, :])
+    latv = 0.5 * (lat[:-1, :] + lat[1:, :])
 
     # TODO wrap up into temporary container Dataset object?
-    return lon1, lat1, lonu, latu, lonv, latv, lone, late
+    return lon, lat, lonu, latu, lonv, latv, lonq, latq
 
 
-def _rotate(lon1, lat1, lonu, latu, lonv, latv, lone, late, rot):
+def _rotate(lon, lat, lonu, latu, lonv, latv, lonq, latq, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees"""
 
-    (lon2, lat2) = rot_sphere(lon1, lat1, rot)
-    (lonu, latu) = rot_sphere(lonu, latu, rot)
-    (lonv, latv) = rot_sphere(lonv, latv, rot)
-    (lone, late) = rot_sphere(lone, late, rot)
+    (lon, lat) = _rot_sphere(lon, lat, rot)
+    (lonu, latu) = _rot_sphere(lonu, latu, rot)
+    (lonv, latv) = _rot_sphere(lonv, latv, rot)
+    (lonq, latq) = _rot_sphere(lonq, latq, rot)
 
-    return lon2, lat2, lonu, latu, lonv, latv, lone, late
+    return lon, lat, lonu, latu, lonv, latv, lonq, latq
 
 
-def _translate(lon2, lat2, lonu, latu, lonv, latv, lone, late, tra_lat, tra_lon):
+def _translate(lon, lat, lonu, latu, lonv, latv, lonq, latq, tra_lat, tra_lon):
     """Translate grid so that the centre lies at the position (tra_lat, tra_lon)"""
 
-    (lon3, lat3) = tra_sphere(lon2, lat2, tra_lat)
-    (lonu, latu) = tra_sphere(lonu, latu, tra_lat)
-    (lonv, latv) = tra_sphere(lonv, latv, tra_lat)
-    (lone, late) = tra_sphere(lone, late, tra_lat)
+    (lon, lat) = _tra_sphere(lon, lat, tra_lat)
+    (lonu, latu) = _tra_sphere(lonu, latu, tra_lat)
+    (lonv, latv) = _tra_sphere(lonv, latv, tra_lat)
+    (lonq, latq) = _tra_sphere(lonq, latq, tra_lat)
 
-    lon4 = lon3 + tra_lon * np.pi / 180
+    lon = lon + tra_lon * np.pi / 180
     lonu = lonu + tra_lon * np.pi / 180
     lonv = lonv + tra_lon * np.pi / 180
-    lone = lone + tra_lon * np.pi / 180
-    lon4[lon4 < -np.pi] = lon4[lon4 < -np.pi] + 2 * np.pi
+    lonq = lonq + tra_lon * np.pi / 180
+
+    lon[lon < -np.pi] = lon[lon < -np.pi] + 2 * np.pi
     lonu[lonu < -np.pi] = lonu[lonu < -np.pi] + 2 * np.pi
     lonv[lonv < -np.pi] = lonv[lonv < -np.pi] + 2 * np.pi
-    lone[lone < -np.pi] = lone[lone < -np.pi] + 2 * np.pi
-    lat4 = lat3
-
-    return lon4, lat4, lonu, latu, lonv, latv, lone, late
+    lonq[lonq < -np.pi] = lonq[lonq < -np.pi] + 2 * np.pi
 
+    return lon, lat, lonu, latu, lonv, latv, lonq, latq
 
-def rot_sphere(lon1, lat1, rot):
 
-    (n, m) = np.shape(lon1)
+def _rot_sphere(lon, lat, rot):
+    (n, m) = np.shape(lon)
+    # convert rotation angle from degrees to radians
     rot = rot * np.pi / 180
 
-    # translate into x,y,z
+    # translate into Cartesian coordinates x,y,z
     # conventions:  (lon,lat) = (0,0)  corresponds to (x,y,z) = ( 0,-r, 0)
     #               (lon,lat) = (0,90) corresponds to (x,y,z) = ( 0, 0, r)
-    x1 = np.sin(lon1) * np.cos(lat1)
-    y1 = np.cos(lon1) * np.cos(lat1)
-    z1 = np.sin(lat1)
+    x1 = np.sin(lon) * np.cos(lat)
+    y1 = np.cos(lon) * np.cos(lat)
+    z1 = np.sin(lat)
 
     # We will rotate these points around the small circle defined by
     # the intersection of the sphere and the plane that
@@ -379,38 +532,33 @@ def rot_sphere(lon1, lat1, rot):
     y2 = y1
     z2 = rp1 * np.sin(ap2)
 
-    lon2 = np.pi / 2 * np.ones((n, m))
-    lon2[abs(y2) > 1e-7] = np.arctan(
+    lon = np.pi / 2 * np.ones((n, m))
+    lon[abs(y2) > 1e-7] = np.arctan(
         np.abs(x2[np.abs(y2) > 1e-7] / y2[np.abs(y2) > 1e-7])
     )
-    lon2[y2 < 0] = np.pi - lon2[y2 < 0]
-    lon2[x2 < 0] = -lon2[x2 < 0]
+    lon[y2 < 0] = np.pi - lon[y2 < 0]
+    lon[x2 < 0] = -lon[x2 < 0]
 
     pr2 = np.sqrt(x2**2 + y2**2)
-    lat2 = np.pi / 2 * np.ones((n, m))
-    lat2[np.abs(pr2) > 1e-7] = np.arctan(
+    lat = np.pi / 2 * np.ones((n, m))
+    lat[np.abs(pr2) > 1e-7] = np.arctan(
         np.abs(z2[np.abs(pr2) > 1e-7] / pr2[np.abs(pr2) > 1e-7])
     )
-    lat2[z2 < 0] = -lat2[z2 < 0]
-
-    return (lon2, lat2)
+    lat[z2 < 0] = -lat[z2 < 0]
 
+    return (lon, lat)
 
-def tra_sphere(lon1, lat1, tra):
 
-    # Rotate sphere around its y-axis
-    # Part of easy grid
-    # (c) 2008, Jeroen Molemaker, UCLA
-
-    (n, m) = np.shape(lon1)
+def _tra_sphere(lon, lat, tra):
+    (n, m) = np.shape(lon)
     tra = tra * np.pi / 180  # translation in latitude direction
 
     # translate into x,y,z
     # conventions:  (lon,lat) = (0,0)  corresponds to (x,y,z) = ( 0,-r, 0)
     #               (lon,lat) = (0,90) corresponds to (x,y,z) = ( 0, 0, r)
-    x1 = np.sin(lon1) * np.cos(lat1)
-    y1 = np.cos(lon1) * np.cos(lat1)
-    z1 = np.sin(lat1)
+    x1 = np.sin(lon) * np.cos(lat)
+    y1 = np.cos(lon) * np.cos(lat)
+    z1 = np.sin(lat)
 
     # We will rotate these points around the small circle defined by
     # the intersection of the sphere and the plane that
@@ -438,29 +586,29 @@ def tra_sphere(lon1, lat1, tra):
     z2 = rp1 * np.sin(ap2)
 
     ## transformation from (x,y,z) to (lat,lon)
-    lon2 = np.pi / 2 * np.ones((n, m))
-    lon2[np.abs(y2) > 1e-7] = np.arctan(
+    lon = np.pi / 2 * np.ones((n, m))
+    lon[np.abs(y2) > 1e-7] = np.arctan(
         np.abs(x2[np.abs(y2) > 1e-7] / y2[np.abs(y2) > 1e-7])
     )
-    lon2[y2 < 0] = np.pi - lon2[y2 < 0]
-    lon2[x2 < 0] = -lon2[x2 < 0]
+    lon[y2 < 0] = np.pi - lon[y2 < 0]
+    lon[x2 < 0] = -lon[x2 < 0]
 
     pr2 = np.sqrt(x2**2 + y2**2)
-    lat2 = np.pi / (2 * np.ones((n, m)))
-    lat2[np.abs(pr2) > 1e-7] = np.arctan(
+    lat = np.pi / (2 * np.ones((n, m)))
+    lat[np.abs(pr2) > 1e-7] = np.arctan(
         np.abs(z2[np.abs(pr2) > 1e-7] / pr2[np.abs(pr2) > 1e-7])
     )
-    lat2[z2 < 0] = -lat2[z2 < 0]
+    lat[z2 < 0] = -lat[z2 < 0]
 
-    return (lon2, lat2)
+    return (lon, lat)
 
 
-def _compute_coordinate_metrics(lon4, lonu, latu, lonv, latv):
+def _compute_coordinate_metrics(lon, lonu, latu, lonv, latv):
     """Compute the curvilinear coordinate metrics pn and pm, defined as 1/grid spacing"""
 
     # pm = 1/dx
     pmu = gc_dist(lonu[:, :-1], latu[:, :-1], lonu[:, 1:], latu[:, 1:])
-    pm = 0 * lon4
+    pm = 0 * lon
     pm[:, 1:-1] = pmu
     pm[:, 0] = pm[:, 1]
     pm[:, -1] = pm[:, -2]
@@ -468,7 +616,7 @@ def _compute_coordinate_metrics(lon4, lonu, latu, lonv, latv):
 
     # pn = 1/dy
     pnv = gc_dist(lonv[:-1, :], latv[:-1, :], lonv[1:, :], latv[1:, :])
-    pn = 0 * lon4
+    pn = 0 * lon
     pn[1:-1, :] = pnv
     pn[0, :] = pn[1, :]
     pn[-1, :] = pn[-2, :]
@@ -478,7 +626,6 @@ def _compute_coordinate_metrics(lon4, lonu, latu, lonv, latv):
 
 
 def gc_dist(lon1, lat1, lon2, lat2):
-
     # Distance between 2 points along a great circle
     # lat and lon in radians!!
     # 2008, Jeroen Molemaker, UCLA
@@ -497,7 +644,7 @@ def gc_dist(lon1, lat1, lon2, lat2):
     return dis
 
 
-def _compute_angle(lon4, lonu, latu, lone):
+def _compute_angle(lon, lonu, latu, lonq):
     """Compute angles of local grid positive x-axis relative to east"""
 
     dellat = latu[:, 1:] - latu[:, :-1]
@@ -506,7 +653,7 @@ def _compute_angle(lon4, lonu, latu, lone):
     dellon[dellon < -np.pi] = dellon[dellon < -np.pi] + 2 * np.pi
     dellon = dellon * np.cos(0.5 * (latu[:, 1:] + latu[:, :-1]))
 
-    ang = copy.copy(lon4)
+    ang = copy.copy(lon)
     ang_s = np.arctan(dellat / (dellon + 1e-16))
     ang_s[(dellon < 0) & (dellat < 0)] = ang_s[(dellon < 0) & (dellat < 0)] - np.pi
     ang_s[(dellon < 0) & (dellat >= 0)] = ang_s[(dellon < 0) & (dellat >= 0)] + np.pi
@@ -517,36 +664,29 @@ def _compute_angle(lon4, lonu, latu, lone):
     ang[:, 0] = ang[:, 1]
     ang[:, -1] = ang[:, -2]
 
-    lon4[lon4 < 0] = lon4[lon4 < 0] + 2 * np.pi
-    lone[lone < 0] = lone[lone < 0] + 2 * np.pi
+    lon[lon < 0] = lon[lon < 0] + 2 * np.pi
+    lonq[lonq < 0] = lonq[lonq < 0] + 2 * np.pi
 
     return ang
 
 
 def _create_grid_ds(
-    nx,
-    ny,
     lon,
     lat,
     pm,
     pn,
     angle,
-    size_x,
-    size_y,
     rot,
     center_lon,
     center_lat,
-    lone,
-    late,
 ):
-
-    # Coriolis frequency
-    f0 = 4 * np.pi * np.sin(lat) / (24 * 3600)
-
-    # Create empty xarray.Dataset object to store variables in
-    ds = xr.Dataset()
-
-    # TODO some of these variables are defined but never written to in Easy Grid
+    # Create xarray.Dataset object with lat_rho and lon_rho as coordinates
+    ds = xr.Dataset(
+        coords={
+            "lat_rho": (("eta_rho", "xi_rho"), lat * 180 / np.pi),
+            "lon_rho": (("eta_rho", "xi_rho"), lon * 180 / np.pi),
+        }
+    )
 
     ds["angle"] = xr.Variable(
         data=angle,
@@ -554,17 +694,16 @@ def _create_grid_ds(
         attrs={"long_name": "Angle between xi axis and east", "units": "radians"},
     )
 
-    # ds['h'] = ...
-    # TODO hraw comes from topography
-    # ds['hraw'] = xr.Variable(data=hraw, dims=['eta_rho', 'xi_rho'])
+    # Coriolis frequency
+    f0 = 4 * np.pi * np.sin(lat) / (24 * 3600)
 
-    ds["f0"] = xr.Variable(
+    ds["f"] = xr.Variable(
         data=f0,
         dims=["eta_rho", "xi_rho"],
         attrs={"long_name": "Coriolis parameter at rho-points", "units": "second-1"},
     )
-    ds["pn"] = xr.Variable(
-        data=pn,
+    ds["pm"] = xr.Variable(
+        data=pm,
         dims=["eta_rho", "xi_rho"],
         attrs={
             "long_name": "Curvilinear coordinate metric in xi-direction",
@@ -572,7 +711,7 @@ def _create_grid_ds(
         },
     )
     ds["pn"] = xr.Variable(
-        data=pm,
+        data=pn,
         dims=["eta_rho", "xi_rho"],
         attrs={
             "long_name": "Curvilinear coordinate metric in eta-direction",
@@ -585,66 +724,86 @@ def _create_grid_ds(
         dims=["eta_rho", "xi_rho"],
         attrs={"long_name": "longitude of rho-points", "units": "degrees East"},
     )
+
     ds["lat_rho"] = xr.Variable(
         data=lat * 180 / np.pi,
         dims=["eta_rho", "xi_rho"],
         attrs={"long_name": "latitude of rho-points", "units": "degrees North"},
     )
 
-    ds["spherical"] = xr.Variable(
-        data=["T"],
-        dims=["one"],
-        attrs={
-            "long_name": "Grid type logical switch",
-            "option_T": "spherical",
-        },
-    )
-
-    # TODO this mask is obtained from hraw
-    # ds['mask_rho'] = xr.Variable(data=lat * 180 / np.pi, dims=['eta_rho', 'xi_rho'], attrs={'long_name': "latitude of rho-points", 'units': "degrees North"})
+    ds["tra_lon"] = center_lon
+    ds["tra_lon"].attrs["long_name"] = "Longitudinal translation of base grid"
+    ds["tra_lon"].attrs["units"] = "degrees East"
 
-    # TODO this 'one' dimension is completely unneccessary as netCDF can store scalars
-    ds["tra_lon"] = xr.Variable(
-        data=[center_lon],
-        dims=["one"],
-        attrs={
-            "long_name": "Longitudinal translation of base grid",
-            "units": "degrees East",
-        },
-    )
-    ds["tra_lat"] = xr.Variable(
-        data=[center_lat],
-        dims=["one"],
-        attrs={
-            "long_name": "Latitudinal translation of base grid",
-            "units": "degrees North",
-        },
-    )
-    ds["rotate"] = xr.Variable(
-        data=[rot],
-        dims=["one"],
-        attrs={"long_name": "Rotation of base grid", "units": "degrees"},
-    )
+    ds["tra_lat"] = center_lat
+    ds["tra_lat"].attrs["long_name"] = "Latitudinal translation of base grid"
+    ds["tra_lat"].attrs["units"] = "degrees North"
 
-    # TODO this is never written to
-    # ds['xy_flip']
+    ds["rotate"] = rot
+    ds["rotate"].attrs["long_name"] = "Rotation of base grid"
+    ds["rotate"].attrs["units"] = "degrees"
 
     return ds
 
 
-def _add_global_metadata(ds, nx, ny, size_x, size_y, center_lon, center_lat, rot):
-
-    ds.attrs["Title"] = (
-        "ROMS grid. Settings:"
-        f" nx: {nx} ny: {ny} "
-        f" xsize: {size_x / 1e3} ysize: {size_y / 1e3}"
-        f" rotate: {rot} Lon: {center_lon} Lat: {center_lat}"
-    )
-    ds.attrs["Date"] = date.today()
+def _add_global_metadata(ds, size_x, size_y):
     ds.attrs["Type"] = "ROMS grid produced by roms-tools"
+    ds.attrs["size_x"] = size_x
+    ds.attrs["size_y"] = size_y
 
     return ds
 
 
-def _make_topography(ds):
-    ...
+def _f2c(f):
+    """
+    Coarsen input xarray DataArray f in both x- and y-direction.
+
+    Parameters
+    ----------
+    f : xarray.DataArray
+        Input DataArray with dimensions (nxp, nyp).
+
+    Returns
+    -------
+    fc : xarray.DataArray
+        Output DataArray with modified dimensions and values.
+    """
+
+    fc = _f2c_xdir(f)
+    fc = fc.transpose()
+    fc = _f2c_xdir(fc)
+    fc = fc.transpose()
+    fc = fc.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
+
+    return fc
+
+
+def _f2c_xdir(f):
+    """
+    Coarsen input xarray DataArray f in x-direction.
+
+    Parameters
+    ----------
+    f : xarray.DataArray
+        Input DataArray with dimensions (nxp, nyp).
+
+    Returns
+    -------
+    fc : xarray.DataArray
+        Output DataArray with modified dimensions and values.
+    """
+    nxp, nyp = f.shape
+    nxcp = (nxp - 2) // 2 + 2
+
+    fc = xr.DataArray(np.zeros((nxcp, nyp)), dims=f.dims)
+
+    # Calculate the interior values
+    fc[1:-1, :] = 0.5 * (f[1:-2:2, :] + f[2:-1:2, :])
+
+    # Calculate the first row
+    fc[0, :] = f[0, :] + 0.5 * (f[0, :] - f[1, :])
+
+    # Calculate the last row
+    fc[-1, :] = f[-1, :] + 0.5 * (f[-1, :] - f[-2, :])
+
+    return fc