#!/usr/bin/env python
# coding: utf-8

# # End-to-end workflow
# 
# <div class="alert alert-info">
# 
# Tip
# 
# Prefer working in a plain Python script instead of a notebook? [Download this tutorial as a .py script](_static/end_to_end.py).
# 
# Prefer the notebook itself? Click the download button at the top of this page and choose `.ipynb` (not `.pdf`).
# </div>
# 
# This notebook demonstrates an end-to-end workflow using `ROMS-Tools` to prepare inputs and analyze outputs for a ROMS-MARBL simulation.
# 
# We will create the following components:
# 
# - Grid  
# - Initial conditions
# - Surface forcing   
# - Boundary forcing  
# - River forcing  
# - CDR (Carbon Dioxide Removal) forcing  
# 
# After generating these inputs, we will run a ROMS-MARBL simulation and perform a basic analysis of the model output using `ROMS-Tools`.
# 
# <div class="alert alert-info">
# 
# Note
# 
# This notebook is intended to provide a high-level overview. For detailed explanations, see the example notebooks that explore each step individually.
# </div>

# ## Running the notebook
# 
# To run this notebook, make sure you have the required dependencies installed. You can set them up in one of the following ways (see also [this installation page](https://roms-tools.readthedocs.io/en/latest/installation.html)):
# 
# 1. **Install from PyPI**  
#    ```bash
#    pip install roms-tools[notebooks]
#    ```
# 2. **Install from GitHub**
#    ```bash
#    cd roms-tools
#    conda env create -f ci/environment-with-xesmf.yml
#    conda activate romstools-test
#    pip install ".[notebooks]"
#    ```
# 3. **Add missing packages to an existing roms-tools installation**
#    ```bash
#    pip install copernicusmarine gdown podman
#    ```
# 
# In addition, `podman` must be installed to run ROMS in a container. See the installation instructions [here](https://podman.io/docs/installation).

# ## Creating the input files with `ROMS-Tools`

# Let's prepare our input data for January 2012.

# In[1]:


from datetime import datetime
start_time = datetime(2012, 1, 1)
end_time = datetime(2012, 1, 31)


# We will save the input datasets to a designated target directory. Feel free to modify the target path!

# In[2]:


from pathlib import Path


# In[3]:


target_dir = Path("./ROMS_TOOLS_INPUT_DATA")


# In[4]:


# Create the directory if it doesn't exist
target_dir.mkdir(exist_ok=True)


# ### Grid

# In[5]:


from roms_tools import Grid


# Let's make a new domain in the Gulf of Mexico, with horizontal resolution of 20km and 20 vertical layers.

# In[6]:


get_ipython().run_cell_magic('time', '', '\ngrid = Grid(\n    nx=100,  # number of grid points in x-direction\n    ny=80,  # number of grid points in y-direction\n    size_x=2000,  # domain size in x-direction (in km)\n    size_y=1600,  # domain size in y-direction (in km)\n    center_lon=-89,  # longitude of the center of the domain\n    center_lat=24,  # latitude of the center of the domain\n    rot=0,  # rotation of the grid (in degrees)\n    N=20,  # number of vertical layers\n)\n')


# <div class="alert alert-info">
# 
# Note
# 
# This example uses the default ETOPO5 bathymetry, with a horizontal resolution of 1/12°. This is sufficient for our 20 km grid, but for higher-resolution grids, consider using SRTM15. For more details on topography source data, see [this page](https://roms-tools.readthedocs.io/en/latest/grid.html#Topography-source-data).
# 
# </div>

# Let’s examine the grid variables that were created, which can be found in the `xarray.Dataset` returned by the `.ds` attribute.

# In[7]:


grid.ds


# To visualize the grid we have just created, we can use the `.plot` method.

# In[8]:


grid.plot()


# We can save our grid variables as a netCDF file via the `.save` method.

# In[9]:


get_ipython().run_cell_magic('time', '', '\ngrid.save(target_dir / "roms_grd.nc")\n')


# We can also export the grid parameters to a YAML file. This gives us a more storage-effective way to save and share input data made with `ROMS-Tools`. The YAML file can be used to recreate the same object later.

# In[10]:


yaml_filepath = target_dir / "roms_grd.yaml"
grid.to_yaml(yaml_filepath)


# These are the contents of the written YAML file.

# In[11]:


with open(yaml_filepath, "r") as file:
    file_contents = file.read()

# Print the contents
print(file_contents)


# You can find more information about creating, plotting, and saving grids [here](https://roms-tools.readthedocs.io/en/latest/grid.html).

# ### Initial Conditions

# We want to prepare a simulation in which we run ROMS with MARBL biogeochemistry (BGC), so we need to prepare both physical and BGC initial conditions. We create physical and BGC initial conditions **together** because ROMS needs a **single** initial conditions file. We use
# 
# * **GLORYS data** to create our physical initial conditions, i.e., initial temperature, salinity, sea surface height, and velocities;
# * a **unified biogeochemical (BGC) climatology** combining multiple observationally and model based sources to create our BGC initial conditions.
# 
# [Here](https://roms-tools.readthedocs.io/en/latest/datasets.html) you can find instruction on how to download these datasets. We have already completed these steps, and the datasets are locally available at the paths shown below.

# In[12]:


glorys_path = "source-data/GoM_GLORYS_Jan2012.nc"


# In[13]:


unified_bgc_path = "source-data/BGCdataset.nc"


# With the source data in place, we can now generate the initial conditions.

# In[14]:


from roms_tools import InitialConditions


# In[15]:


get_ipython().run_cell_magic('time', '', '\ninitial_conditions = InitialConditions(\n    grid=grid,\n    ini_time=start_time,\n    source={"name": "GLORYS", "path": glorys_path},\n    bgc_source={\n        "name": "UNIFIED",\n        "path": unified_bgc_path,\n        "climatology": True,\n    },\n    use_dask=True\n)\n')


# <div class="alert alert-info">
# 
# Note
# 
# In the cell above, we created our initial conditions with `use_dask = True`. This enables `Dask`, a Python library designated to facilitate scalable, out-of-memory data processing by distributing computations across multiple threads or processes. [Here](https://roms-tools.readthedocs.io/en/latest/using_dask.html) you can learn more about using `Dask` with `ROMS-Tools`.
# 
# </div>

# The following variables were generated as part of the initial conditions.

# In[16]:


initial_conditions.ds


# Here is a plot of the uppermost layer of the temperature field.

# In[17]:


initial_conditions.plot("temp", s=-1)


# We can also plot vertical sections of our initial conditions, for example alkalinity.

# In[18]:


initial_conditions.plot("ALK", eta=50)


# We can now save the dataset as a NetCDF and YAML file.

# In[19]:


initial_conditions.save(target_dir / "roms_ini.nc")


# In[20]:


initial_conditions.to_yaml(target_dir / "roms_ini.yaml")


# You can find more information about creating, plotting, and saving initial conditions [here](https://roms-tools.readthedocs.io/en/latest/initial_conditions.html).

# ### Surface Forcing
# Next, we create the surface forcing for our domain.

# `ROMS-Tools` can create two types of surface forcing:
# 
# * **physical surface forcing**, e.g., 10m wind, shortwave radiation, and air temperature at 2m;
# * **biogeochemical (BGC) surface forcing**, e.g., atmospheric pCO2.
# 
# Unlike initial conditions data, ROMS can read multiple surface forcing files, so we create these two types separately.

# In[21]:


from roms_tools import SurfaceForcing


# #### Physical Surface Forcing

# In this subsection, we use ERA5 data to generate the physical surface forcing. With ROMS-Tools, ERA5 can be streamed directly from the cloud, so you don’t need to pre-download the source data or provide a local path (though you can still do so if you prefer). See [here](https://roms-tools.readthedocs.io/en/latest/surface_forcing.html#Physical-surface-forcing) for more details on source data options.

# In[22]:


get_ipython().run_cell_magic('time', '', '\nsurface_forcing = SurfaceForcing(\n    grid=grid,\n    start_time=start_time,\n    end_time=end_time,\n    source={"name": "ERA5"},\n    type="physics",\n    use_dask=True\n)\n')


# <div class="alert alert-info">
# 
# Note
# 
# For an explanation of the info message about interpolation onto the fine versus coarse grid, see this [page](https://roms-tools.readthedocs.io/en/latest/surface_forcing.html#Fine-vs.-coarse-grid).
# 
# </div>

# These are the physical surface forcing variables that were created.

# In[23]:


surface_forcing.ds


# Here is a plot of the zonal wind at the first created time stamp.

# In[24]:


surface_forcing.plot("uwnd", time=0)


# Let's save our physical surface forcing to NetCDF and YAML files.

# In[25]:


surface_forcing.save(target_dir / "roms_frc.nc")


# In[26]:


surface_forcing.to_yaml(target_dir / "roms_frc.yaml")


# #### Biogeochemical (BGC) Surface Forcing
# We now create BGC surface forcing.

# In[27]:


get_ipython().run_cell_magic('time', '', '\nbgc_surface_forcing = SurfaceForcing(\n    grid=grid,\n    start_time=start_time,\n    end_time=end_time,\n    source={"name": "UNIFIED", "path": unified_bgc_path, "climatology": True},\n    type="bgc",\n    use_dask=True\n)\n')


# Here are the BGC surface variables that were created.

# In[28]:


bgc_surface_forcing.ds


# Let's plot iron, which is one of the BGC surface forcings.

# In[29]:


bgc_surface_forcing.plot("iron", time=0)


# ##### Time-varying CO2 
# 
# By default, ROMS assumes a constant value for the molar mixing ratio for CO2, xco2. If we would like a time-varying xco2 applied in the model, we can create a BGC surface forcing object containing xco2, similar to the BGC surface forcing generated above.  
# 
# The xco2 data is interpolated from the GML's (NOAA) Marine Boundary Layer dataset (1979-2025). For more information, see [details on surface forcing](https://roms-tools.readthedocs.io/en/latest/surface_forcing.html). Here we let the code assume defaults.

# In[30]:


get_ipython().run_cell_magic('time', '', '\nmbl_bgc_surface_forcing = SurfaceForcing(\n    grid=grid,\n    start_time=start_time,\n    end_time=end_time,\n    source={"name": "MBL_co2"},\n    type="bgc",\n)\n')


# Here we plot a time step of xco2.

# In[31]:


mbl_bgc_surface_forcing.plot("xco2_air", time=0)


# Finally, we save our BGC surface forcing to NetCDF and YAML files.

# In[32]:


bgc_surface_forcing.save(target_dir /"roms_frc_bgc.nc")


# In[33]:


bgc_surface_forcing.to_yaml(target_dir / "roms_frc_bgc.yaml")


# You can find more information about creating, plotting, and saving surface forcing [here](https://roms-tools.readthedocs.io/en/latest/surface_forcing.html).

# ### Boundary Forcing
# In this subsection, we create the boundary forcing (or, equivalently, the open boundary conditions).

# `ROMS-Tools` can create two types of boundary forcing:
# 
# * **physical boundary forcing** like temperature, salinity, velocities, and sea surface height;
# * **biogeochemical (BGC) boundary forcing** like alkalinity, dissolved inorganic phosphate, etc.
# 
# As with surface forcing, ROMS accepts multiple boundary forcing files, so we create these two types separately.

# In[34]:


from roms_tools import BoundaryForcing


# #### Physical boundary forcing
# 
# We use the same GLORYS data as above to create our physical boundary forcing. Since the western boundary of our Gulf of Mexico domain lies entirely on land, we can exclude the western boundary in the following.

# In[35]:


get_ipython().run_cell_magic('time', '', '\nboundary_forcing = BoundaryForcing(\n    grid=grid,\n    boundaries={\n        "south": True,\n        "east": True,\n        "north": True,\n        "west": False,   # western boundary excluded\n    },\n    start_time=start_time,\n    end_time=end_time,\n    source={"name": "GLORYS", "path": glorys_path},\n    type="physics",\n    use_dask=True\n)\n')


# <div class="alert alert-info">
# 
# Note
# 
# Boundary data is regridded near land with a NaN-aware approach (masked bilinear interpolation plus nearest-neighbor extrapolation). See this [page](https://roms-tools.readthedocs.io/en/latest/boundary_forcing.html#Horizontal-regridding-near-land-and-coarse-source-data) for details.
# 
# </div>

# Here are the physical boundary forcing variables that were created.

# In[36]:


boundary_forcing.ds


# Let's plot meridional velocity at the southern boundary for the first available time stamp.

# In[37]:


boundary_forcing.plot("u_east", time=0)


# Finally, we save our physical boundary forcing variables to NetCDF and YAML files.

# In[38]:


boundary_forcing.save(target_dir / "roms_bry.nc")


# In[39]:


boundary_forcing.to_yaml(target_dir / "roms_bry.yaml")


# #### Biogeochemical (BGC) boundary forcing
# We now create BGC boundary forcing. The BGC variables are interpolated from the same unified BGC climatology as above.

# In[40]:


get_ipython().run_cell_magic('time', '', '\nbgc_boundary_forcing = BoundaryForcing(\n    grid=grid,\n    boundaries={\n        "south": True,\n        "east": True,\n        "north": True,\n        "west": False,   # western boundary excluded\n    },\n    start_time=start_time,\n    end_time=end_time,\n    source={"name": "UNIFIED", "path": unified_bgc_path, "climatology": True},\n    type="bgc",\n    use_dask=True\n)\n')


# In[41]:


bgc_boundary_forcing.plot("ALK_east", time=0)


# In[42]:


bgc_boundary_forcing.save(target_dir / "roms_bry_bgc.nc")


# In[43]:


bgc_boundary_forcing.to_yaml(target_dir / "roms_bry_bgc.yaml")


# You can find more information about creating, plotting, and saving boundary forcing [here](https://roms-tools.readthedocs.io/en/latest/boundary_forcing.html).

# ### River Forcing
# 
# In this subsection, we create the river forcing.

# In[44]:


from roms_tools import RiverForcing


# In[45]:


get_ipython().run_line_magic('time', '')
river_forcing = RiverForcing(
    grid=grid,
    start_time=start_time,
    end_time=end_time,
    include_bgc=True
)


# Here are the river forcing variables that were created.

# In[46]:


river_forcing.ds


# Let's make a plot of the river locations.

# In[47]:


river_forcing.plot_locations()


# The left plot shows the original river locations from the Dai and Trenberth global river dataset mapped onto the ROMS domain, many of which are off the coast. ROMS requires rivers to be on land adjacent to wet points. The right plot shows the updated locations, where `ROMS-Tools` shifted each river to the nearest coastal grid point.

# Let's make a plot of the volume flux for the rivers shown above.

# In[48]:


river_forcing.plot("river_volume")


# Finally, we save the river forcing to NetCDF and YAML files.

# In[49]:


river_forcing.save(target_dir / "roms_rivers.nc")


# In[50]:


river_forcing.to_yaml(target_dir / "roms_rivers.yaml")


# You can find more information about creating, plotting, and saving river forcing [here](https://roms-tools.readthedocs.io/en/latest/river_forcing.html).

# ### Carbon Dioxide Removal (CDR) Forcing
# Lastly, we create CDR forcing for our simulation.

# In[51]:


from roms_tools import TracerPerturbation, CDRForcing


# We first define a tracer perturbation release, which perturbs the BGC tracer fields without adding water. This approach is suitable for large-scale applications where mixing has already occurred, such as in a pre-run mixing model.
# 
# The perturbation is centered at 28°N, 96°W, and at the surface (0 m depth). It is distributed using a Gaussian with standard deviations of 100 km horizontally and 50 m vertically. We release an alkalinity flux of 2,000,000 meq/s over the entire simulation period.

# In[52]:


release = TracerPerturbation(
    name="OAE",
    lat=28,  # degree N
    lon=-96,  # degree E
    depth=0,  # m
    hsc=100000, # m
    vsc=50,  # m
    tracer_fluxes={"ALK": 2 * 10**6},  # meq/s
)


# For more complex, time-varying releases, or releases that add water in addition to perturbing tracers (relevant for field-scale deployments where mixing is still required), see this [page](https://roms-tools.readthedocs.io/en/latest/cdr_forcing.html).
# 
# We now assemble all releases that we want to run in a single simulations into the CDR forcing.

# In[53]:


cdr_forcing = CDRForcing(
    grid=grid,
    start_time=start_time, # simulation start time
    end_time=end_time, # simulation end time
    releases=[release]
)


# We can make a plot of the release distribution.

# In[54]:


cdr_forcing.plot_distribution(release_name="OAE")


# Here is a plot of the alkalinity flux being released, which is constant over time.

# In[55]:


cdr_forcing.plot_tracer_flux(tracer_name="ALK")


# Finally, we save the CDR forcing to NetCDF and YAML files.

# In[56]:


cdr_forcing.save(target_dir / "roms_cdr.nc")


# In[57]:


cdr_forcing.to_yaml(target_dir / "roms_cdr.yaml")


# ## Running ROMS
# 
# In the last section we have made all necessary input fields. We are now ready to run a ROMS-MARBL simulation!
# 
# Here we run a pre-compiled version of ROMS using a container. Containers can be used to run programs in portable, reproducible pre-configured environments.
# 
# <div class="alert alert-info">
# 
# Note
# 
# This precompiled ROMS executable is built specifically for this example. As a result, compile-time options, such as the horizontal grid dimensions (`nx`, `ny`) used in this notebook, cannot be modified and still be expected to work with the precompiled ROMS executable in the container.
# 
# </div>

# Let's first make a directory that will hold the model output.

# In[58]:


from pathlib import Path


# In[59]:


output_dir = Path("./OUTPUT")


# In[60]:


output_dir.mkdir(exist_ok=True)


# Next, users on MacOS must first initialise a virtual machine. Make sure that `podman` is installed, as noted at the top of the notebook.

# In[61]:


import os


# In[62]:


os.environ["PATH"] = "/opt/podman/bin:" + os.environ["PATH"]


# In[63]:


get_ipython().run_cell_magic('bash', '', '\npodman machine init --cpus 10\npodman machine start\n')


# We have built a [container for this example](https://github.com/orgs/CWorthy-ocean/packages/container/package/roms_tools_end_to_end_example) that is available through the GitHub Container Repository. Let's pull this container image!

# In[64]:


get_ipython().run_cell_magic('bash', '', '\npodman pull ghcr.io/cworthy-ocean/roms_tools_end_to_end_example:latest > /dev/null 2>&1\n')


# In the above cell, the `> /dev/null 2>&1` part suppresses all output. Remove it if you want to see the download progress and logs.

# Now we can run the container with the following flags:
# 
# * `--rm`: automatically remove the container after it finishes
# * `-v`: mount directories for input files and output so ROMS-MARBL can access inputs and save results
# 
# We prefix the command with `time` to measure how long the simulation takes to run.

# In[ ]:


get_ipython().run_cell_magic('bash', '', '\npodman run --rm -v ./ROMS_TOOLS_INPUT_DATA:/input -v ./OUTPUT:/output ghcr.io/cworthy-ocean/roms_tools_end_to_end_example:latest > /dev/null 2>&1\n')


# Again, the `> /dev/null 2>&1` part suppresses all output. Remove it if you want to see the full (and very long) ROMS model output.

# ## Analyzing the output
# The output is now in the `OUTPUT` directory that we created.

# In[ ]:


ls OUTPUT/


# The `ROMS_example_his*.nc` files contain snapshots of the physical variables, while the `ROMS_example_cstar*.nc` files hold the CDR-related output.
# 
# We can visualize the output with `ROMS-Tools`.

# In[ ]:


from roms_tools import ROMSOutput


# Let's first look at the snapshots of the physical variables.

# In[ ]:


roms_output = ROMSOutput(
    grid=grid,
    path="OUTPUT/ROMS_example_his.20120101000000.nc",
    use_dask=True,
)


# This output contains 31 snapshots, one per day.

# In[ ]:


roms_output.ds


# Let’s plot the sea surface temperature at the start of the simulation (January 1, 2012) and again 31 days later.

# In[ ]:


roms_output.plot("temp", s=-1, time=0)


# In[ ]:


roms_output.plot("temp", s=-1, time=-1)


# We can also compare zonal surface velocity at the start of the simulation and 31 days later.

# In[ ]:


roms_output.plot("u", s=-1, time=0)


# In[ ]:


roms_output.plot("u", s=-1, time=-1)


# You can learn more about visualizing ROMS output [here](https://roms-tools.readthedocs.io/en/latest/plotting_roms_output.html).

# Let's now have a look at the CDR-related output.

# In[ ]:


cdr_roms_output = ROMSOutput(
    grid=grid,
    path="OUTPUT/ROMS_example_cstar.20120102000000.nc",
    use_dask=True,
)


# This output contains 30 time stamps, each representing a daily average.

# In[ ]:


cdr_roms_output.ds


# Here we show the alkalinity source term diagnosed online by the model, plotted at the surface for the first time step.

# In[ ]:


cdr_roms_output.plot("ALK_source", s=-1, time=0)


# Finally, we compute the CDR uptake curves using two approaches: from CO₂ flux differences and from DIC differences. A more detailed explanation is provided [here](https://roms-tools.readthedocs.io/en/latest/cdr_analysis.html).

# In[ ]:


cdr_roms_output.cdr_metrics()


# There are slight differences between the uptake curves depending on the computation method, but overall they agree closely. Since the simulation was run for only one month, the uptake efficiency remains relatively low, around 5% at the end of the period. Extending the simulation would allow the efficiency to increase further.
# 
# The computed CDR metrics can also be found in the `.ds_cdr` attribute.

# In[ ]:


cdr_roms_output.ds_cdr


# Learn more about analyzing CDR metrics [here](https://roms-tools.readthedocs.io/en/latest/cdr_analysis.html).

# In[ ]: