update plot scripts for PALEOmodel v0.16

PALEOtoolkit · Jan 14, 2025 · f7a558a · f7a558a
1 parent e1e46ce
commit f7a558a
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 42 deletions.
diff --git a/examples/mitgcm/MITgcm_ECCO_PO4MMbase.jl b/examples/mitgcm/MITgcm_ECCO_PO4MMbase.jl
@@ -56,7 +56,7 @@ end
 initial_state, modeldata = PALEOmodel.initialize!(model; method_barrier) 
 
 for iseg in 1:num_segments
-# for iseg in 9:num_segments
+# for iseg in 20:num_segments
 # iseg = 1
 # if true
     @info """
@@ -119,9 +119,9 @@ println()
 gr(size=(1200, 900))
 pager = PALEOmodel.PlotPager((2,2), (legend_background_color=nothing, margin=(5, :mm)))
 
-plot_forcings(paleorun.output, pager=pager, lonidx=200)
+plot_forcings(paleorun.output, pager=pager, lon=200.0)
 pager(:newpage)
-plot_abiotic_O2(paleorun.output, toutputs=toutputs, pager=pager, lonidx1=200, lonidx2=340)
+plot_abiotic_O2(paleorun.output, toutputs=toutputs, pager=pager, lon1=200.0, lon2=340.0)
 pager(:newpage)
 plot_PO4MMbase(
     paleorun.output,

diff --git a/examples/mitgcm/Test xarray netcdf.ipynb b/examples/mitgcm/Test xarray netcdf.ipynb
@@ -11,12 +11,14 @@
     "Generate file 'test.nc' with:\n",
     "\n",
     "    julia> include(\"MITgcm_2deg8_PO4MMcarbSCH4.jl\")\n",
-    "    julia> PALEOmodel.OutputWriters.save_netcdf(paleorun.output, \"MITgcm_PO4MMcarbSCH42deg8.nc\")"
+    "    julia> PALEOmodel.OutputWriters.save_netcdf(paleorun.output, \"MITgcm_PO4MMcarbSCH42deg8.nc\")\n",
+    "\n",
+    "Uses xarray v2024.10.0 for [xarray.DataTree](https://xarray.dev/blog/datatree) (DataTree is convenient but not required, see commented-out alternate code for earlier xarray versions without DataTree)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "id": "58320e6c",
    "metadata": {},
    "outputs": [],
@@ -27,7 +29,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "d640c5a0",
    "metadata": {},
    "outputs": [],
@@ -37,7 +39,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "c9a1bd15",
    "metadata": {},
    "outputs": [],
@@ -60,26 +62,52 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "401e5aef",
+   "id": "eb755113",
    "metadata": {},
    "outputs": [],
    "source": [
+    "# read whole grouped netcdf file as an xarray datatree\n",
+    "dt = xr.open_datatree(output_filename, decode_cf=True, decode_coords=\"all\")\n",
+    "\n",
     "# read top-level dataset\n",
-    "ds = xr.open_dataset(output_filename, decode_cf=True)\n",
-    "ds.attrs"
+    "# ds = xr.open_dataset(output_filename, decode_cf=True)\n",
+    "\n",
+    "dt.attrs"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
-   "id": "4b80a015",
+   "execution_count": null,
+   "id": "3a062196",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# read ocean group (= PALEO Domain ocean) from netcdf file\n",
-    "ds_ocean = xr.open_dataset(output_filename, group=\"ocean\", decode_cf=True, decode_coords=\"all\")\n",
+    "dt.children"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "bddbdc01",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get ocean group (= PALEO Domain ocean) from datatree\n",
+    "ds_ocean = dt[\"ocean\"].to_dataset()\n",
     "num_records = len(ds_ocean.tmodel)\n",
     "\n",
+    "# read ocean group (= PALEO Domain ocean) from netcdf file\n",
+    "# ds_ocean = xr.open_dataset(output_filename, group=\"ocean\", decode_cf=True, decode_coords=\"all\")\n",
+    "# num_records = len(ds_ocean.tmodel)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "4b80a015",
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "# attach z coordinates (this is not automatic!)\n",
     "# ds_ocean = ds_ocean.set_coords([\"zmid\", \"zlower\", \"zupper\"])"
    ]
@@ -329,7 +357,7 @@
    "source": [
     "# columns from 3 lon sections at lat = 0\n",
     "H2S_conc_depth_3lon = ds_ocean.H2S_conc.isel(tmodel=num_records-1).sel(isotopelinear=\"total\").sel(lon=[70.0, 200.0, 260.0, 340.0], lat=0, method=\"nearest\") \n",
-    "H2S_conc_depth_3lon.plot.line(y=\"zmid\", hue=\"lon\")\n",
+    "H2S_conc_depth_3lon.plot.line(y=\"zt\", hue=\"lon\")\n",
     "# TODO - use zlower, zupper to create a stepped plot"
    ]
   }
@@ -350,7 +378,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.14"
+   "version": "3.10.16"
   }
  },
  "nbformat": 4,

diff --git a/examples/mitgcm/images/make_images.jl b/examples/mitgcm/images/make_images.jl
@@ -1,58 +1,59 @@
 
-# gr(size=(1200, 900))
+using Plots
+using Printf
 
+"""
+    make_images(output)
 
-using Printf
+Produce single-panel svg figures for ./images
 
-# produce figures for ./images
+Defaults and titles are for `output` produced by `examples/mitgcm/MITgcm_2deg8_PO4MMcarbSCH4.jl`
+"""
 function make_images(
     output;
-    lonidx1=71, # 2.8 Pac   200 ECCO
-    lonidx2=121, # 2.8 Atl  340 ECCO
+    lon1=198.28, # Pac cell centre for 2.8 deg
+    lon2=338.91, # Atl cell centre for 2.8 deg
+    toutput = 1e12, # last timestep
+    show_sections=true,
+    lon_lims = (0.0, 360.0),
+    lat_lims = (-90.0, 90.0),
 )
-    toutput = 1e12 # last timestep
     tmodel = last(PALEOmodel.get_array(output, "ocean.tmodel").values)
     println("using tmodel $tmodel (yr)")
 
-
-    # get lon, lat coords from surface P
-    P_conc = PALEOmodel.get_array(output, "ocean.P_conc", (k=1, tmodel=toutput))
-    lon, lon_lower, lon_upper = P_conc.dims[1].coords
-    lat, lat_lower, lat_upper = P_conc.dims[2].coords
-    lon1 = lon.values[lonidx1]
-    lon2 = lon.values[lonidx2]
-    println("longitudes for sections $lonidx1 = $lon1 (deg), $lonidx2 = $lon2 (deg)")    
-    # lat_lims = (first(lat_lower.values), last(lat_upper.values))
-    lat_lims = (-90.0, 90.0)
-    println("lat limits $lat_lims")
-
+    println("longitudes for sections $lon1 (deg), $lon2 (deg)")
+
+    println("lon limits $lon_lims, lat limits $lat_lims")
 
     gr(size=(700, 500))
     p = heatmap(
-        title=@sprintf("Surface P (mol m-3) at %.2f yr", tmodel), output, "ocean.P_conc", (tmodel=toutput, k=1); 
-        swap_xy=true, xlabel="lon (deg)", ylabel="lat (deg)", margin=(5, :mm), ylims=lat_lims,
+        title=@sprintf("Surface P (mol m-3) at %.2f yr", tmodel), 
+        output, "ocean.P_conc", (tmodel=toutput, zt_isel=1, expand_cartesian=true); 
+        swap_xy=true, xlabel="lon (deg)", ylabel="lat (deg)", margin=(5, :mm), xlims=lon_lims, ylims=lat_lims,
     )
-    plot!(p, [lon1, lon1],  [-90, 90], linecolor=:red, linewidth=2, label=nothing)
-    plot!(p, [lon2, lon2],  [-90, 90], linecolor=:red, linewidth=2, label=nothing)
+    if show_sections
+        plot!(p, [lon1, lon1],  [-90, 90], linecolor=:red, linewidth=2, label=nothing)
+        plot!(p, [lon2, lon2],  [-90, 90], linecolor=:red, linewidth=2, label=nothing)
+    end
     display(p)
     savefig(p, "surface_P_2deg8_PO4MMcarbSCH4_2000yr.svg")
 
     gr(size=(800, 400))
     p = heatmap(
-        title=@sprintf("H2S (mol m-3) lon %.2f at %.2f yr", lon1, tmodel), output, "ocean.H2S_conc", (tmodel=toutput, i=lonidx1);
+        title=@sprintf("H2S (mol m-3) lon %.2f at %.2f yr", lon1, tmodel),
+        output, "ocean.H2S_conc", (tmodel=toutput, lon=lon1, expand_cartesian=true);
         swap_xy=true, xlabel="lat (deg)", ylabel="depth (m)", clims=(0.0, 0.11), margin=(5, :mm),
     )
     display(p)
     savefig(p, "H2S_lon1_2deg8_PO4MMcarbSCH4_2000yr.svg")
 
     p = heatmap(
-        title=@sprintf("H2S (mol m-3) lon %.2f at %.2f yr", lon2, tmodel), output, "ocean.H2S_conc", (tmodel=toutput, i=lonidx2);
+        title=@sprintf("H2S (mol m-3) lon %.2f at %.2f yr", lon2, tmodel), 
+        output, "ocean.H2S_conc", (tmodel=toutput, lon=lon2, expand_cartesian=true);
         swap_xy=true, xlabel="lat (deg)", ylabel="depth (m)", clims=(0.0, 0.11), margin=(5, :mm),
     ) 
     display(p)
     savefig(p, "H2S_lon2_2deg8_PO4MMcarbSCH4_2000yr.svg")
 
-    # heatmap(output, "ocean.$tr", (tmodel=t, i=lonidx2), swap_xy=true), 
-
     return nothing
 end