Add new feature ‘custom diagnostics’ to Fluid2D

core/default.xml

@@ -224,6 +224,11 @@
       <default>all</default>
       <doc>List of variables names that we want to output in the diagnostic file. The variables names have to be consistent with the model. If 'all', store them all.</doc>
     </name>
+    <name value="custom_diag">
+      <type>str</type>
+      <default>none</default>
+      <doc>Dictionary of custom diagnostics that are added to the diagnostic file. The keys are the names under which the diagnostics are saved. The values of the dictionary are functions that are called at every time diagnostics are evaluated. The functions take as only argument the current state of the model and return a single value (float).</doc>
+    </name>
     <name value="nprint">
       <type>int,long</type>
       <default>20</default>

core/defaults.json

@@ -246,6 +246,11 @@
             "default": "all",
             "doc": "List of variables names that we want to output in the diagnostic file. The variables names have to be consistent with the model. If 'all', store them all."
         },
+        "custom_diag": {
+            "type": "str",
+            "default": "none",
+            "doc": "Dictionary of custom diagnostics that are added to the diagnostic file. The keys are the names under which the diagnostics are saved. The values of the dictionary are functions that are called at every time diagnostics are evaluated. The functions take as only argument the current state of the model and return a single value (float)."
+        },
         "nprint": {
             "type": "int",
             "default": 20,

core/output.py

@@ -12,7 +12,7 @@ def __init__(self, param, grid, diag, flxlist=None):
                            'var_to_save', 'varname_list',
                            'expdir', 'tracer_list',
                            'diag_fluxes',
-                           'freq_his', 'freq_diag', 'list_diag']
+                           'freq_his', 'freq_diag', 'list_diag', 'custom_diag']
         param.copy(self, self.list_param)
 
         self.list_grid = ['nxl', 'nyl', 'nh']
@@ -50,6 +50,8 @@ def __init__(self, param, grid, diag, flxlist=None):
         # prepare the 'diagnostics' file
         if self.list_diag == 'all':
             self.list_diag = diag.keys()
+        if self.custom_diag == 'none':
+            self.custom_diag = {}
 
         self.diagfile = '%s/%s_diag.nc' % (self.expdir, self.expname)
 
@@ -85,7 +87,7 @@ def do(self, data, t, kt):
         if (t >= self.tnextdiag):
             self.tnextdiag += self.freq_diag
             if self.myrank == 0:
-                self.write_diag(data['diag'], t, kt)
+                self.write_diag(data['his'], data['diag'], t, kt)
 
         if (t >= self.tnexthis):
             self.tnexthis += self.freq_his
@@ -112,16 +114,19 @@ def create_diag(self, diag):
             for v in self.list_diag:
                 d = nc.createVariable(v, 'f', ('t',))
                 d.long_name = v
+            for v in self.custom_diag:
+                d = nc.createVariable(v, 'f', ('t',))
+                d.long_name = v
 
         self.kdiag = 0
 
         # set up internal buffer to avoid too frequent disk access
-        self.ndiags = len(self.list_diag)+2
+        self.ndiags = len(self.list_diag) + len(self.custom_diag) + 2
         self.buffersize = 10
         self.buffer = np.zeros((self.buffersize, self.ndiags))
 
 
-    def write_diag(self, diag, t, kt):
+    def write_diag(self, his, diag, t, kt):
 
         # store diag into the buffer
         k = self.kdiag % self.buffersize
@@ -131,6 +136,10 @@ def write_diag(self, diag, t, kt):
         for v in self.list_diag:
             self.buffer[k, j] = diag[v]  # getattr(diag,v)
             j += 1
+        for v in self.custom_diag:
+            # Call the custom diagnostics function with the current model state
+            self.buffer[k, j] = self.custom_diag[v](his)
+            j += 1
         self.kdiag += 1
 
         # write buffer into netcdf if full
@@ -149,6 +158,9 @@ def dump_diag(self):
         for v in self.list_diag:
             nc.variables[v][k] = self.buffer[:last, j]
             j += 1
+        for v in self.custom_diag:
+            nc.variables[v][k] = self.buffer[:last, j]
+            j += 1
         nc.close()
 
     def join(self):

experiments/Vortex/vortex_on_beta_plane.py

@@ -0,0 +1,133 @@
+"""QG-simulation of a vortex on the beta-plane.
+
+This experiment provides an example for the usage of custom diagnostics.
+Custom diagnostics are used to get the position of the vortex with a
+high temporal resolution (here: daily).  This can then be used to
+calculate accurately the velocity of the vortex due to the beta-effect.
+
+To make additional information appear in the diagnostics file, we use
+the parameter custom_diag which we set to a dictionary as in the example
+below.  The keys of the dictionary are strings which represent the names
+under which the information is stored in the NetCDF file.  The values
+are functions -- typical Python style.  They take as argument the
+current state of the model and return a single value, usually a float.
+
+While the same information could be extracted later from the history
+file, the advantage of using custom diagnostics is that we get a good
+resolution in time with a small size of the output files.  For this, we
+choose a high frequency output to write diagnostics, so we get the
+information we want with very good temporal resolution, and we choose a
+low frequency output for the history file, so that its file size does
+not become too big.
+
+This is extremely useful in bulk studies with the Experiment Management
+System of Fluid2D, which makes it easy to run hundreds of experiments
+automatically with small changes in parameters.
+
+"""
+
+from fluid2d import Fluid2d
+from param import Param
+from grid import Grid
+
+import numpy as np
+
+
+# Define custom diagnostic functions
+def get_vortex_x(state):
+    """Calculate the x-position of the center of a vortex.
+
+    The implementation is for a vortex with positive relative vorticity
+    and uses the global variables ‘param’ and ‘grid’.
+    """
+    pvanom = state[param.varname_list.index("pvanom")]
+    # Take the maximum along y, then get the argmax along x and return
+    # the corresponding coordinate
+    return grid.x1d[np.argmax(np.max(pvanom, axis=0))]
+
+def get_vortex_y(state):
+    """Calculate the y-position of the center of a vortex.
+
+    The implementation is for a vortex with positive relative vorticity
+    and uses the global variables ‘param’ and ‘grid’.
+    """
+    pvanom = state[param.varname_list.index("pvanom")]
+    return grid.y1d[np.argmax(np.max(pvanom, axis=1))]
+
+
+param = Param("default.xml")
+
+# Choose a name for the output of the experiment
+param.expname = "QG_vortex"
+
+# Set the type of model and domain, its size and resolution
+param.modelname = "quasigeostrophic"
+param.geometry = "xchannel"
+param.nx = 64 * 2
+param.ny = 64 * 2
+# We can think of lengths as being in kilometer
+param.Lx = 1000
+param.Ly = 1000
+# Set the number of CPU cores used
+param.npx = 1
+param.npy = 1
+
+# Configure physics
+param.forcing = False
+param.noslip = False
+param.diffusion = False
+param.Rd = 100
+param.beta = 0.001
+
+# Set time settings; we can think of one time unit as one day
+param.tend = 365
+param.adaptable_dt = True
+# for adaptable timestepping:
+param.cfl = 0.8
+param.dtmax = 1
+# for fixed timestepping:
+param.dt = 0.5
+
+# Set the discretization
+param.order = 5
+param.timestepping = "RK3_SSP"
+param.exacthistime = True
+
+# Configure the output with our own diagnostics
+param.var_to_save = ["psi", "pv", "pvanom"]
+param.list_diag = "all"
+param.custom_diag = {"x_pos": get_vortex_x, "y_pos": get_vortex_y}
+# Set a high frequency for ‘diagnostics’ (every day) to have a high
+# temporal resolution; set a low frequency for ‘history’ (every 10 days)
+# to keep the filesize small
+param.freq_his = 10
+param.freq_diag = 1
+
+# Set plot settings
+param.plot_interactive = True
+param.generate_mp4 = False  # requires plot_interactive to be True
+param.freq_plot = 20
+param.plot_psi = False
+param.plot_var = "pv"
+param.cmap = "plasma"
+
+# Create model
+grid = Grid(param)
+f2d = Fluid2d(param, grid)
+model = f2d.model
+
+# Set the initial vorticity field
+pv = model.var.get("pv")
+# Get centered coordinates
+x = grid.xr - param.Lx / 2
+y = grid.yr - param.Ly / 2
+# Create a vortex that is n-times stronger than the vorticity background
+n = 2
+amplitude = n * param.beta * param.Ly
+width = 50
+pv[:] = model.pvback + amplitude * np.exp(-(x**2 + y**2) / width**2)
+model.ope.fill_halo(pv[:])
+model.set_psi_from_pv()
+
+# Start simulation
+f2d.loop()