Merge pull request #269 from OceanParcels/brownian_diffusion

Brownian diffusion

Author: delandmeterp

docs/index.rst

@@ -33,6 +33,7 @@ on the `Github Development Timeline page
 **Major features**
 
 * Advection of particles in 2D using inbuilt kernels for Runge-Kutta4, Runge-Kutta45 and Euler Forward and in 3D using the inbuilt kernel for Runge-Kutta4_3D (see :mod:`parcels.kernels.advection`)
+* Simple horizontal diffusion of particles using inbuilt Brownian Motion kernel (see :mod:`parcels.kernels.diffusion`)
 * Ability to define and execute custom kernels (see `the Adding-a-custom-behaviour-kernel part of the Tutorial <http://nbviewer.jupyter.org/github/OceanParcels/parcels/blob/master/parcels/examples/parcels_tutorial.ipynb#Adding-a-custom-behaviour-kernel>`_)
 * Ability to add custom Variables to Particles (see `the Sampling-a-Field-with-Particles part of the Tutorial <http://nbviewer.jupyter.org/github/OceanParcels/parcels/blob/master/parcels/examples/parcels_tutorial.ipynb#Sampling-a-Field-with-Particles>`_)
 * Ability to add and remove Particles (see :func:`parcels.particleset.ParticleSet.add` and :func:`parcels.particleset.ParticleSet.remove`)
@@ -49,7 +50,7 @@ on the `Github Development Timeline page
 
 **Future development goals**
 
-* Diffusion of particles using suite of inbuilt kernels
+* More types of diffusion of particles using suite of inbuilt kernels
 * Support for unstructured grids
 * Implementation of parallel execution using tiling of the domain
 * Faster and more efficient code

parcels/examples/example_brownian.py

@@ -1,65 +1,59 @@
-from parcels import FieldSet, ParticleSet, ScipyParticle, JITParticle
+from parcels import FieldSet, Field, ParticleSet, ScipyParticle, JITParticle, BrownianMotion2D
 import numpy as np
 from datetime import timedelta as delta
-import math
 from parcels import random
 import pytest
 
 ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
 
 
-def two_dim_brownian_flat(particle, fieldset, time, dt):
-    # Kernel for simple Brownian particle diffusion in zonal and meridional direction.
-
-    particle.lat += random.normalvariate(0, 1)*math.sqrt(2*dt*fieldset.Kh_meridional)
-    particle.lon += random.normalvariate(0, 1)*math.sqrt(2*dt*fieldset.Kh_zonal)
-
-
-def brownian_fieldset(xdim=200, ydim=200):     # Define a flat fieldset of zeros, for simplicity.
-    dimensions = {'lon': np.linspace(0, 600000, xdim, dtype=np.float32),
-                  'lat': np.linspace(0, 600000, ydim, dtype=np.float32)}
+def zeros_fieldset(xdim=2, ydim=2):
+    """Generates a zero velocity field"""
+    lon = np.linspace(-20, 20, xdim, dtype=np.float32)
+    lat = np.linspace(-20, 20, ydim, dtype=np.float32)
 
+    dimensions = {'lon': lon, 'lat': lat}
     data = {'U': np.zeros((xdim, ydim), dtype=np.float32),
             'V': np.zeros((xdim, ydim), dtype=np.float32)}
-
-    return FieldSet.from_data(data, dimensions, mesh='flat')
+    return FieldSet.from_data(data, dimensions, mesh='spherical')
 
 
 @pytest.mark.parametrize('mode', ['scipy', 'jit'])
 def test_brownian_example(mode, npart=3000):
-    fieldset = brownian_fieldset()
+    fieldset = zeros_fieldset()
 
     # Set diffusion constants.
-    fieldset.Kh_meridional = 100.
-    fieldset.Kh_zonal = 100.
+    kh_zonal = 100
+    kh_meridional = 100
+
+    # Create field of Kh_zonal and Kh_meridional, using same grid as U
+    grid = fieldset.U.grid
+    fieldset.add_field(Field('Kh_zonal', kh_zonal*np.ones((2, 2)), grid=grid))
+    fieldset.add_field(Field('Kh_meridional', kh_meridional*np.ones((2, 2)), grid=grid))
 
     # Set random seed
     random.seed(123456)
 
-    ptcls_start = 300000.  # Start all particles at same location in middle of grid.
-    pset = ParticleSet.from_line(fieldset=fieldset, size=npart, pclass=ptype[mode],
-                                 start=(ptcls_start, ptcls_start),
-                                 finish=(ptcls_start, ptcls_start))
-
-    endtime = delta(days=1)
-    dt = delta(hours=1)
-    interval = delta(hours=1)
+    runtime = delta(days=1)
 
-    k_brownian = pset.Kernel(two_dim_brownian_flat)
+    random.seed(1234)
+    pset = ParticleSet(fieldset=fieldset, pclass=ptype[mode],
+                       lon=np.zeros(npart), lat=np.zeros(npart))
+    pset.execute(pset.Kernel(BrownianMotion2D),
+                 runtime=runtime, dt=delta(hours=1))
 
-    pset.execute(k_brownian, endtime=endtime, dt=dt, interval=interval,
-                 output_file=pset.ParticleFile(name="BrownianParticle"),
-                 show_movie=False)
+    expected_std_x = np.sqrt(2*kh_zonal*runtime.total_seconds())
+    expected_std_y = np.sqrt(2*kh_meridional*runtime.total_seconds())
 
-    lats = np.array([particle.lat for particle in pset.particles])
-    lons = np.array([particle.lon for particle in pset.particles])
-    expected_std_lat = np.sqrt(2*fieldset.Kh_meridional*endtime.total_seconds())
-    expected_std_lon = np.sqrt(2*fieldset.Kh_zonal*endtime.total_seconds())
+    conversion = (1852 * 60)  # to convert from degrees to m
+    ys = np.array([p.lat for p in pset]) * conversion
+    xs = np.array([p.lon for p in pset]) * conversion  # since near equator, we do not need to care about curvature effect
 
-    assert np.allclose(np.std(lats), expected_std_lat, rtol=.1)
-    assert np.allclose(np.std(lons), expected_std_lon, rtol=.1)
-    assert np.allclose(np.mean(lons), ptcls_start, rtol=.1)
-    assert np.allclose(np.mean(lats), ptcls_start, rtol=.1)
+    tol = 200  # 200m tolerance
+    assert np.allclose(np.std(xs), expected_std_x, atol=tol)
+    assert np.allclose(np.std(ys), expected_std_y, atol=tol)
+    assert np.allclose(np.mean(xs), 0, atol=tol)
+    assert np.allclose(np.mean(ys), 0, atol=tol)
 
 
 if __name__ == "__main__":

parcels/field.py

@@ -14,7 +14,7 @@
                   CurvilinearSGrid, CGrid, GridCode)
 
 
-__all__ = ['CentralDifferences', 'Field', 'Geographic', 'GeographicPolar']
+__all__ = ['CentralDifferences', 'Field', 'Geographic', 'GeographicPolar', 'GeographicSquare', 'GeographicPolarSquare']
 
 
 class FieldSamplingError(RuntimeError):
@@ -139,6 +139,44 @@ def ccode_to_source(self, x, y, z):
         return "(1000. * 1.852 * 60. * cos(%s * M_PI / 180))" % y
 
 
+class GeographicSquare(UnitConverter):
+    """ Square distance converter from geometric to geographic coordinates (m2 to degree2) """
+    source_unit = 'm2'
+    target_unit = 'degree2'
+
+    def to_target(self, value, x, y, z):
+        return value / pow(1000. * 1.852 * 60., 2)
+
+    def to_source(self, value, x, y, z):
+        return value * pow(1000. * 1.852 * 60., 2)
+
+    def ccode_to_target(self, x, y, z):
+        return "pow(1.0 / (1000.0 * 1.852 * 60.0), 2)"
+
+    def ccode_to_source(self, x, y, z):
+        return "pow((1000.0 * 1.852 * 60.0), 2)"
+
+
+class GeographicPolarSquare(UnitConverter):
+    """ Square distance converter from geometric to geographic coordinates (m2 to degree2)
+        with a correction to account for narrower grid cells closer to the poles.
+    """
+    source_unit = 'm2'
+    target_unit = 'degree2'
+
+    def to_target(self, value, x, y, z):
+        return value / pow(1000. * 1.852 * 60. * cos(y * pi / 180), 2)
+
+    def to_source(self, value, x, y, z):
+        return value * pow(1000. * 1.852 * 60. * cos(y * pi / 180), 2)
+
+    def ccode_to_target(self, x, y, z):
+        return "pow(1.0 / (1000. * 1.852 * 60. * cos(%s * M_PI / 180)), 2)" % y
+
+    def ccode_to_source(self, x, y, z):
+        return "pow((1000. * 1.852 * 60. * cos(%s * M_PI / 180)), 2)" % y
+
+
 class Field(object):
     """Class that encapsulates access to field data.
 
@@ -168,6 +206,10 @@ class Field(object):
            This flag overrides the allow_time_interpolation and sets it to False
     """
 
+    unitconverters = {'U': GeographicPolar(), 'V': Geographic(),
+                      'Kh_zonal': GeographicPolarSquare(),
+                      'Kh_meridional': GeographicSquare()}
+
     def __init__(self, name, data, lon=None, lat=None, depth=None, time=None, grid=None,
                  transpose=False, vmin=None, vmax=None, time_origin=0,
                  interp_method='linear', allow_time_extrapolation=None, time_periodic=False, mesh='flat'):
@@ -186,12 +228,10 @@ def __init__(self, name, data, lon=None, lat=None, depth=None, time=None, grid=N
         self.lat = self.grid.lat
         self.depth = self.grid.depth
         self.time = self.grid.time
-        if self.grid.mesh is 'flat' or (name is not 'U' and name is not 'V'):
+        if self.grid.mesh is 'flat' or (name not in self.unitconverters.keys()):
             self.units = UnitConverter()
-        elif self.grid.mesh is 'spherical' and name == 'U':
-            self.units = GeographicPolar()
-        elif self.grid.mesh is 'spherical' and name == 'V':
-            self.units = Geographic()
+        elif self.grid.mesh is 'spherical':
+            self.units = self.unitconverters[name]
         else:
             raise ValueError("Unsupported mesh type. Choose either: 'spherical' or 'flat'")
         self.interp_method = interp_method

parcels/kernels/__init__.py

@@ -1,2 +1,3 @@
 from .advection import *  # noqa
+from .diffusion import *  # noqa
 from .error import *  # noqa

parcels/kernels/diffusion.py

@@ -0,0 +1,16 @@
+from parcels import rng as random
+import math
+
+
+__all__ = ['BrownianMotion2D']
+
+
+def BrownianMotion2D(particle, fieldset, time, dt):
+    # Kernel for simple Brownian particle diffusion in zonal and meridional direction.
+    # Assumes that fieldset has fields Kh_zonal and Kh_meridional
+
+    r = 1/3.
+    kh_meridional = fieldset.Kh_meridional[time, particle.lon, particle.lat, particle.depth]
+    particle.lat += random.uniform(-1., 1.) * math.sqrt(2*math.fabs(dt)*kh_meridional/r)
+    kh_zonal = fieldset.Kh_zonal[time, particle.lon, particle.lat, particle.depth]
+    particle.lon += random.uniform(-1., 1.) * math.sqrt(2*math.fabs(dt)*kh_zonal/r)

tests/test_diffusion.py

@@ -0,0 +1,53 @@
+from parcels import (FieldSet, Field, RectilinearZGrid, ParticleSet, BrownianMotion2D,
+                     JITParticle, ScipyParticle)
+from parcels import rng as random
+from datetime import timedelta as delta
+import numpy as np
+import pytest
+
+ptype = {'scipy': ScipyParticle, 'jit': JITParticle}
+
+
+def zeros_fieldset(mesh='spherical', xdim=200, ydim=100, mesh_conversion=1):
+    """Generates a zero velocity field"""
+    lon = np.linspace(-1e5*mesh_conversion, 1e5*mesh_conversion, xdim, dtype=np.float32)
+    lat = np.linspace(-1e5*mesh_conversion, 1e5*mesh_conversion, ydim, dtype=np.float32)
+
+    dimensions = {'lon': lon, 'lat': lat}
+    data = {'U': np.zeros((xdim, ydim), dtype=np.float32),
+            'V': np.zeros((xdim, ydim), dtype=np.float32)}
+    return FieldSet.from_data(data, dimensions, mesh=mesh)
+
+
+@pytest.mark.parametrize('mesh', ['spherical', 'flat'])
+@pytest.mark.parametrize('mode', ['scipy', 'jit'])
+def test_fieldKh_Brownian(mesh, mode, xdim=200, ydim=100, kh_zonal=100, kh_meridional=50):
+    mesh_conversion = 1/1852./60 if mesh is 'spherical' else 1
+    fieldset = zeros_fieldset(mesh=mesh, xdim=xdim, ydim=ydim, mesh_conversion=mesh_conversion)
+
+    vec = np.linspace(-1e5*mesh_conversion, 1e5*mesh_conversion, 2)
+    grid = RectilinearZGrid('K_grid', lon=vec, lat=vec, mesh=mesh)
+
+    fieldset.add_field(Field('Kh_zonal', kh_zonal*np.ones((2, 2)), grid=grid))
+    fieldset.add_field(Field('Kh_meridional', kh_meridional*np.ones((2, 2)), grid=grid))
+
+    npart = 1000
+    runtime = delta(days=1)
+
+    random.seed(1234)
+    pset = ParticleSet(fieldset=fieldset, pclass=ptype[mode],
+                       lon=np.zeros(npart), lat=np.zeros(npart))
+    pset.execute(pset.Kernel(BrownianMotion2D),
+                 runtime=runtime, dt=delta(hours=1))
+
+    expected_std_lon = np.sqrt(2*kh_zonal*mesh_conversion**2*runtime.total_seconds())
+    expected_std_lat = np.sqrt(2*kh_meridional*mesh_conversion**2*runtime.total_seconds())
+
+    lats = np.array([p.lat for p in pset])
+    lons = np.array([p.lon for p in pset])
+
+    tol = 200*mesh_conversion  # effectively 200 m errors
+    assert np.allclose(np.std(lats), expected_std_lat, atol=tol)
+    assert np.allclose(np.std(lons), expected_std_lon, atol=tol)
+    assert np.allclose(np.mean(lons), 0, atol=tol)
+    assert np.allclose(np.mean(lats), 0, atol=tol)